diff options
| -rw-r--r-- | data/__init__.py | 6 |
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 |
