(dataset: Dataset, batch_size: int)
| 74 | |
| 75 | |
| 76 | def prepare_dataloader(dataset: Dataset, batch_size: int): |
| 77 | return DataLoader( |
| 78 | dataset, |
| 79 | batch_size=batch_size, |
| 80 | pin_memory=True, |
| 81 | shuffle=False, |
| 82 | sampler=DistributedSampler(dataset) |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int): |