Update log

pytorch · Jan 29, 2025 · aab9239 · aab9239
1 parent 438bc60
commit aab9239
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/torchft/checkpointing.py b/torchft/checkpointing.py
@@ -134,9 +134,9 @@ def do_GET(self):
                         self.end_headers()
 
                         state_dict = ckpt_server._state_dict
-                        self._logger.warning("Before torch.save ===================.")
+                        logger.warning("Before torch.save ===================.")
                         torch.save(state_dict, self.wfile)
-                        self._logger.warning("After torch.save ===================.")
+                        logger.warning("After torch.save ===================.")
 
                 except Exception as e:
                     logger.exception(

diff --git a/torchft/manager.py b/torchft/manager.py
@@ -608,12 +608,12 @@ def load_state_dict(self, state_dict: Dict[str, int]) -> None:
         self._batches_committed = state_dict["batches_committed"]
 
     def _manager_state_dict(self) -> Dict[str, object]:
-        self._logger.warning("Before state_dict ===================.")
+        self._logger.warn("Before state_dict ===================.")
         ret = {
             "user": self._user_state_dict(),
             "torchft": self.state_dict(),
         }
-        self._logger.warning("After state_dict ===================.")
+        self._logger.warn("After state_dict ===================.")
         return ret
 
     def state_dict(self) -> Dict[str, int]: