diff --git a/python/hdf5.py b/python/hdf5.py
index 4acbccfb3acca59fbff54e052f831442218d156b..13681011d82fa99e7a21f4e533dbc75592aa4521 100644
--- a/python/hdf5.py
+++ b/python/hdf5.py
@@ -142,20 +142,19 @@ if __name__ == '__main__':
 
 #%%
 import h5py
-def hdf5Generator(filePath, batch_size, dataSet):
+def hdf5Generator(filePath, batch_size, dataSet,loop=True):
     with h5py.File(filePath, 'r') as hf:
         L = len(hf["X" + dataSet])
         while True:
             batch_start = 0
             batch_end = batch_size
             
-            while batch_start < L:
-                limit = min(batch_end, L)
-                X = hf["X" + dataSet][batch_start:limit]
-                Y = hf["Y" + dataSet][batch_start:limit]
-    
+            while batch_end < L:
+                X = hf["X" + dataSet][batch_start:batch_end]
+                Y = hf["Y" + dataSet][batch_start:batch_end]
                 yield (X,Y) #a tuple with two numpy arrays with batch_size samples     
     
                 batch_start += batch_size   
                 batch_end += batch_size
+            if not loop: break
    
\ No newline at end of file