From 27524ad0856bb776750d3cd1d0770496f08f2452 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 4 Jul 2026 06:35:29 +0800
Subject: [PATCH] fix: reset sampler iter at epoch end so progress bar shows
 total after first epoch

---
 astrai/dataset/sampler.py | 1 +
 1 file changed, 1 insertion(+)

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):