(dataset: Dataset, batch_size: int)
| 82 | |
| 83 | |
| 84 | def prepare_dataloader(dataset: Dataset, batch_size: int): |
| 85 | return DataLoader( |
| 86 | dataset, |
| 87 | batch_size=batch_size, |
| 88 | pin_memory=True, |
| 89 | shuffle=False, |
| 90 | sampler=DistributedSampler(dataset) |
| 91 | ) |
| 92 | |
| 93 | |
| 94 | def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"): |