add infinite dataloader

utils/utils.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import numpy as np
+from torch.utils.data import DataLoader
 
 
 
@@ -25,3 +26,20 @@ class Accumulator():
         self.dict = dict()
 
 
+class InfiniteDataLoader(DataLoader):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Initialize an iterator over the dataset.
+        self.dataset_iterator = super().__iter__()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            batch = next(self.dataset_iterator)
+        except StopIteration:
+            # Dataset exhausted, use a new fresh iterator.
+            self.dataset_iterator = super().__iter__()
+            batch = next(self.dataset_iterator)
+        return batch