Merge pull request #147 from SsnL/max_ds

Fix max_dataset_size

1 files changed, 7 insertions, 1 deletions

diff --git a/data/custom_dataset_data_loader.py b/data/custom_dataset_data_loader.py
index 60180e0..787946f 100644
--- a/data/custom_dataset_data_loader.py
+++ b/data/custom_dataset_data_loader.py
@@ -35,7 +35,13 @@ class CustomDatasetDataLoader(BaseDataLoader):
             num_workers=int(opt.nThreads))
 
     def load_data(self):
-        return self.dataloader
+        return self
 
     def __len__(self):
         return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        for i, data in enumerate(self.dataloader):
+            if i >= self.opt.max_dataset_size:
+                break
+            yield data