update

Author: deependujha

src/lightning/pytorch/strategies/strategy.py

@@ -363,7 +363,9 @@ def lightning_module(self) -> Optional["pl.LightningModule"]:
         """Returns the pure LightningModule without potential wrappers."""
         return self._lightning_module
 
-    def load_checkpoint(self, checkpoint_path: _PATH, weights_only: Optional[bool] = None, state: dict[str, Any] = None) -> dict[str, Any]:
+    def load_checkpoint(
+        self, checkpoint_path: _PATH, weights_only: Optional[bool] = None, state: dict[str, Any] = None
+    ) -> dict[str, Any]:
         torch.cuda.empty_cache()
         return self.checkpoint_io.load_checkpoint(checkpoint_path, weights_only=weights_only, state=state)
 

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -81,7 +81,9 @@ def resume_start(self, checkpoint_path: Optional[_PATH] = None, weights_only: Op
         rank_zero_info(f"Restoring states from the checkpoint path at {checkpoint_path}")
         with pl_legacy_patch():
             _current_state = self.dump_checkpoint(weights_only=weights_only)
-            loaded_checkpoint = self.trainer.strategy.load_checkpoint(checkpoint_path, weights_only=weights_only, state=_current_state)
+            loaded_checkpoint = self.trainer.strategy.load_checkpoint(
+                checkpoint_path, weights_only=weights_only, state=_current_state
+            )
         self._loaded_checkpoint = _pl_migrate_checkpoint(loaded_checkpoint, checkpoint_path)
 
     def _select_ckpt_path(

src/lightning/pytorch/trainer/trainer.py

@@ -1063,7 +1063,7 @@ def _run(
 
         if self.strategy.restore_checkpoint_after_setup:
             log.debug(f"{self.__class__.__name__}: restoring module and callbacks from checkpoint path: {ckpt_path}")
-            self._checkpoint_connector._restore_modules_and_callbacks(ckpt_path)
+            self._checkpoint_connector._restore_modules_and_callbacks(ckpt_path, weights_only)
 
         # restore optimizers, etc.
         log.debug(f"{self.__class__.__name__}: restoring training state")