From e3b79532e4eb1b5a32cb5a43a180e196996fc86d Mon Sep 17 00:00:00 2001
From: Jon Freer <j.freer@hotmail.co.uk>
Date: Tue, 17 Nov 2020 19:24:10 +0000
Subject: [PATCH] epoch fix

---
 utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/utils.py b/utils.py
index bdca336c..07d3a572 100644
--- a/utils.py
+++ b/utils.py
@@ -491,7 +491,7 @@ def __init__(self, data_source, num_epochs, start_itr=0, batch_size=128):
     self.data_source = data_source
     self.num_samples = len(self.data_source)
     self.num_epochs = num_epochs
-    self.start_itr = start_itr
+    self.start_itr = start_itr%int((self.num_samples*self.num_epochs)/batch_size) 
     self.batch_size = batch_size
 
     if not isinstance(self.num_samples, int) or self.num_samples <= 0:
@@ -514,6 +514,7 @@ def __iter__(self):
     # return iter(.tolist())
     output = torch.cat(out).tolist()
     print('Length dataset output is %d' % len(output))
+    self.start_itr = 0
     return iter(output)
 
   def __len__(self):