diff options
Diffstat (limited to 'data/__init__.py')
| -rw-r--r-- | data/__init__.py | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/data/__init__.py b/data/__init__.py index ef581e7..a69f374 100644 --- a/data/__init__.py +++ b/data/__init__.py @@ -20,6 +20,9 @@ def CreateDataset(opt): elif opt.dataset_mode == 'single': from data.single_dataset import SingleDataset dataset = SingleDataset() + elif opt.dataset_mode == 'recursive': + from data.recursive_dataset import RecursiveDataset + dataset = RecursiveDataset() else: raise ValueError("Dataset [%s] not recognized." % opt.dataset_mode) @@ -49,6 +52,6 @@ class CustomDatasetDataLoader(BaseDataLoader): def __iter__(self): for i, data in enumerate(self.dataloader): - if i >= self.opt.max_dataset_size: + if i * self.opt.batchSize >= self.opt.max_dataset_size: break yield data |