fix(trainer): 修复检查点加载逻辑

2025-10-18 21:45:23 +08:00
parent b67bc9865d
commit 622982364b
5 changed files with 35 additions and 21 deletions
@@ -273,14 +273,15 @@ class ResumeableRandomSampler(Sampler[int]):
        
        generator = torch.Generator()
        generator.manual_seed(seed)
+
+        # consume  previous epochs
+        for _ in range(start_epoch):
+            torch.randperm(self.num_samples, generator=generator)
        
        self.generator = generator
-        self._indices = None 
+        self._indices = None
    
    def _get_indices(self):
-        for _ in range(self.epoch):
-            _ = torch.randperm(self.num_samples, generator=self.generator)
-        
        current_epoch_indices = torch.randperm(self.num_samples, generator=self.generator).tolist()
        self._indices = current_epoch_indices[self.iter % self.num_samples:]