add comment

pytorch · Feb 12, 2025 · da55b34 · da55b34
1 parent 53f37a0
commit da55b34
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchdata/stateful_dataloader/sampler.py b/torchdata/stateful_dataloader/sampler.py
@@ -69,8 +69,8 @@ def load_state_dict(self, state_dict: dict) -> None:
         self.sampler.generator.set_state(self.generator_state)
 
         if self.next_yielded is not None:
-            self.perm = self._get_perm()
-            for _ in range(self.next_yielded - self.yielded):
+            self.perm = self._get_perm()  # We want permutations from the latest generator state that's loaded
+            for _ in range(self.next_yielded):
                 next(self)
             self.yielded = self.next_yielded
             self.next_yielded = None