Implement distributed data parallel via torch elastic launcher

1 files changed, 11 insertions, 0 deletions

diff --git a/libs/datautils.py b/libs/datautils.py
index 6a7c506..53222a8 100644
--- a/libs/datautils.py
+++ b/libs/datautils.py
@@ -125,3 +125,14 @@ class TwinTransform:
         v1 = self.transform(x)
         v2 = self.transform(x)
         return v1, v2
+
+
+class ContinuousSampler(torch.utils.data.sampler.Sampler):
+    def __init__(self, sampler):
+        super(ContinuousSampler, self).__init__(sampler)
+        self.base_sampler = sampler
+
+    def __iter__(self):
+        while True:
+            for batch in self.base_sampler:
+                yield batch