From 957a2a46e7725184776c3c72860e8215164cc4ef Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Sat, 13 Aug 2022 23:38:43 +0800
Subject: Implement distributed data parallel via torch elastic launcher

---
 libs/datautils.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

(limited to 'libs/datautils.py')

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
-- 
cgit v1.2.3