updated train.py

IshanRattan · Aug 7, 2024 · 49b0abe · 49b0abe
1 parent 5fabafe
commit 49b0abe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dl/image/classification/cifar10/train.py b/dl/image/classification/cifar10/train.py
@@ -15,7 +15,7 @@
 
 
 datasets = {'train':train_set, 'val':val_set}
-dataloaders = {x: torch.utils.data.DataLoader(dataset=datasets[x], batch_size=batch_size, shuffle=True)
+dataloaders = {x: torch.utils.data.DataLoader(dataset=datasets[x], batch_size=config.batch_size, shuffle=True)
               for x in ['train', 'val']}
 
 dataset_sizes = {x: len(datasets[x]) for x in ['train', 'val']}