Update SampleGenerator

torch_xla/utils/utils.py

@@ -8,6 +8,8 @@
 import tempfile
 import time
 
+from torch.utils.data.dataloader import DataLoader
+
 
 class Cleaner(object):
 
@@ -38,7 +40,7 @@ def __init__(self):
     self.cleaner = Cleaner(lambda: shutil.rmtree(self.name))
 
 
-class SampleGenerator(object):
+class SampleGenerator(DataLoader):
   """Iterator which returns multiple samples of a given input data.
 
   Can be used in place of a PyTorch `DataLoader` to generate synthetic data.
@@ -54,7 +56,7 @@ def __init__(self, data, sample_count):
     self._count = 0
 
   def __iter__(self):
-    return SampleGenerator(self._data, self._sample_count)
+    return self
 
   def __len__(self):
     return self._sample_count