From 63761c5d8abea169d8a03a73c0bae99aaf036b6a Mon Sep 17 00:00:00 2001 From: junyanz Date: Thu, 20 Apr 2017 06:35:55 -0700 Subject: fix a bug in unaligned_data_loader --- data/unaligned_data_loader.py | 4 ++++ 1 file changed, 4 insertions(+) (limited to 'data') diff --git a/data/unaligned_data_loader.py b/data/unaligned_data_loader.py index 4f82dbe..4a06510 100644 --- a/data/unaligned_data_loader.py +++ b/data/unaligned_data_loader.py @@ -13,6 +13,8 @@ class PairedData(object): self.stop_B = False def __iter__(self): + self.stop_A = False + self.stop_B = False self.data_loader_A_iter = iter(self.data_loader_A) self.data_loader_B_iter = iter(self.data_loader_B) return self @@ -37,6 +39,8 @@ class PairedData(object): B, B_paths = next(self.data_loader_B_iter) if self.stop_A and self.stop_B: + self.stop_A = False + self.stop_B = False raise StopIteration() else: return {'A': A, 'A_paths': A_paths, -- cgit v1.2.3-70-g09d2