diff --git a/astrai/dataset/sampler.py b/astrai/dataset/sampler.py
index 2d6bd0c..6d857b8 100644
--- a/astrai/dataset/sampler.py
+++ b/astrai/dataset/sampler.py
@@ -74,6 +74,7 @@ class ResumableDistributedSampler(Sampler[int]):
 
         self.epoch += 1
         self._indices = None
+        self.iter = self.iter % self.num_samples_per_replica
 
     @property
     def _remaining(self):