From 1cf7171641711d69663733989c267110d66ca466 Mon Sep 17 00:00:00 2001 From: junyanz Date: Sun, 15 Apr 2018 00:47:21 -0400 Subject: fix a bug (batchSize >1) in the data_loader --- data/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'data') 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 -- cgit v1.2.3-70-g09d2