From 59fb33ba65e0714926304aec444c92a55ab03a49 Mon Sep 17 00:00:00 2001 From: SsnL Date: Thu, 9 Nov 2017 17:41:51 -0500 Subject: fix max_dataset_size --- data/custom_dataset_data_loader.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) (limited to 'data/custom_dataset_data_loader.py') 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 -- cgit v1.2.3-70-g09d2