run precommit

pytorch · Feb 7, 2025 · 4de1bb4 · 4de1bb4
1 parent 1346e44
commit 4de1bb4
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 166 deletions.
diff --git a/torchdata/stateful_dataloader/sampler.py b/torchdata/stateful_dataloader/sampler.py
@@ -109,18 +109,15 @@ 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) -> None:
-        if isinstance(self.sampler_iter, Stateful) and hasattr(
-            self.sampler_iter, "update_state_dict"
-        ):
+        if isinstance(self.sampler_iter, Stateful) and hasattr(self.sampler_iter, "update_state_dict"):
             self.sampler_iter.update_state_dict()