diff options
Diffstat (limited to 'libs/datautils.py')
-rw-r--r-- | libs/datautils.py | 11 |
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 |