diff options
| -rw-r--r-- | data/__init__.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/data/__init__.py b/data/__init__.py index ef581e7..341281d 100644 --- a/data/__init__.py +++ b/data/__init__.py @@ -49,6 +49,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 |