typing

Author: fegin

torchft/manager.py

@@ -534,11 +534,11 @@ def _apply_pending_state_dict(self) -> None:
 
         self._logger.info("applying pending state dict")
 
+        assert self._pending_state_dict is not None, "checkpoint was not staged"
         assert (
             self._load_state_dict is not None
         ), "user load_state_dict is not initialized."
         self._load_state_dict(self._pending_state_dict["user"])
-        assert self._pending_state_dict is not None, "checkpoint was not staged"
         self._pending_state_dict = None
         self._logger.info("Loaded state dict.")