summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--data/__init__.py6
1 files changed, 5 insertions, 1 deletions
diff --git a/data/__init__.py b/data/__init__.py
index 381ec59..6da79f3 100644
--- a/data/__init__.py
+++ b/data/__init__.py
@@ -70,7 +70,11 @@ class RecursiveDatasetDataLoader(BaseDataLoader):
#BaseDataLoader.initialize(self, opt)
self.opt = opt
self.dataset = CreateDataset(opt)
- self.dataloader = self.dataset
+ self.dataloader = torch.utils.data.DataLoader(
+ self.dataset,
+ batch_size=opt.batchSize,
+ shuffle=not opt.serial_batches,
+ num_workers=int(opt.nThreads))
def load_data(self):
return self