fix: reset sampler iter at epoch end so progress bar shows total after first epoch

2026-07-04 06:35:29 +08:00 · 2026-07-04 06:35:29 +08:00 · 27524ad085
parent 27d1921d9c
commit 27524ad085
1 changed files with 1 additions and 0 deletions
--- 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):