add comment about why were updating state dict

pytorch · Feb 7, 2025 · 1346e44 · 1346e44
1 parent 8a83c5f
commit 1346e44
Show file tree

Hide file tree

Showing 2 changed files with 170 additions and 58 deletions.
diff --git a/torchdata/stateful_dataloader/sampler.py b/torchdata/stateful_dataloader/sampler.py
@@ -109,15 +109,18 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             assert isinstance(self.sampler_iter, Stateful)
             self.sampler_iter.load_state_dict(state_dict[self._SAMPLER_ITER_STATE])
 
-        if not (isinstance(self.sampler, Stateful) or isinstance(self.sampler_iter, Stateful)) and not isinstance(
-            self.sampler, _InfiniteConstantSampler
-        ):
+        if not (
+            isinstance(self.sampler, Stateful)
+            or isinstance(self.sampler_iter, Stateful)
+        ) and not isinstance(self.sampler, _InfiniteConstantSampler):
             # We skip x samples if underlying sampler is not stateful
             for _ in range(self.samples_yielded):
                 next(self.sampler_iter)
 
-    def update_state_dict(self):
-        if isinstance(self.sampler_iter, Stateful) and hasattr(self.sampler_iter, "update_state_dict"):
+    def update_state_dict(self) -> None:
+        if isinstance(self.sampler_iter, Stateful) and hasattr(
+            self.sampler_iter, "update_state_dict"
+        ):
             self.sampler_iter.update_state_dict()