add num_threads to train.py

richzhang · Jun 4, 2020 · 66a1cb2 · 66a1cb2
1 parent 1f7575e
commit 66a1cb2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -26,7 +26,7 @@
                                                                             transforms.RandomResizedCrop(opt.fineSize, interpolation=3)]),
                                                    transforms.RandomHorizontalFlip(),
                                                    transforms.ToTensor()]))
-    dataset_loader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True)
+    dataset_loader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=int(opt.num_threads))
 
     dataset_size = len(dataset)
     print('#training images = %d' % dataset_size)