define methods in _StatefulDistributedSamplerIterator

pytorch · Aug 21, 2024 · 0dce976 · 0dce976
1 parent fe08bfc
commit 0dce976
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/torchdata/stateful_dataloader/sampler.py b/torchdata/stateful_dataloader/sampler.py
@@ -41,7 +41,6 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         self.generator_state = state_dict[self._GENERATOR]
         self.sampler.generator.set_state(state_dict[self._GENERATOR])
         self.next_yielded = state_dict[self._YIELDED]
-        return None
 
     def state_dict(self) -> Dict[str, Any]:
         return {self._GENERATOR: self.generator_state, self._YIELDED: self.yielded}
@@ -175,6 +174,12 @@ def __next__(self) -> int:
         self.sampler.yielded += 1
         return val
 
+    def state_dict(self) -> Dict[str, Any]:
+        return self.sampler.state_dict()
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        self.sampler.load_state_dict(state_dict)
+
 
 class StatefulDistributedSampler(torch.utils.data.distributed.DistributedSampler):
     _YIELDED = "yielded"