(dataset: Dataset, batch_size: int)
| 55 | |
| 56 | |
| 57 | def prepare_dataloader(dataset: Dataset, batch_size: int): |
| 58 | return DataLoader( |
| 59 | dataset, |
| 60 | batch_size=batch_size, |
| 61 | pin_memory=True, |
| 62 | shuffle=True |
| 63 | ) |
| 64 | |
| 65 | |
| 66 | def main(device, total_epochs, save_every, batch_size): |