weights_only as last arg

Author: matsumotosan

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

@@ -230,7 +230,7 @@ def resume_end(self) -> None:
         # wait for all to catch up
         self.trainer.strategy.barrier("_CheckpointConnector.resume_end")
 
-    def restore(self, checkpoint_path: Optional[_PATH] = None) -> None:
+    def restore(self, checkpoint_path: Optional[_PATH] = None, weights_only: Optional[bool] = None) -> None:
         """Attempt to restore everything at once from a 'PyTorch-Lightning checkpoint' file through file-read and
         state-restore, in this priority:
 
@@ -244,7 +244,7 @@ def restore(self, checkpoint_path: Optional[_PATH] = None) -> None:
             checkpoint_path: Path to a PyTorch Lightning checkpoint file.
 
         """
-        self.resume_start(checkpoint_path)
+        self.resume_start(checkpoint_path, weights_only=weights_only)
 
         # restore module states
         self.restore_datamodule()

src/lightning/pytorch/trainer/trainer.py

@@ -557,6 +557,12 @@ def fit(
                         - ``'registry:version:v2'``: uses the default model set
                           with ``Trainer(..., model_registry="my-model")`` and version 'v2'
 
+            weights_only: Defaults to ``None``. If ``True``, restricts loading to ``state_dicts`` of plain
+                ``torch.Tensor`` and other primitive types. If loading a checkpoint from a trusted source that contains
+                an ``nn.Module``, use ``weights_only=False``. If loading checkpoint from an untrusted source, we
+                recommend using ``weights_only=True``. For more information, please refer to the
+                `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
+
         Raises:
             TypeError:
                 If ``model`` is not :class:`~lightning.pytorch.core.LightningModule` for torch version less than
@@ -630,9 +636,9 @@ def validate(
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
+        weights_only: Optional[bool] = None,
     ) -> _EVALUATE_OUTPUT:
         r"""Perform one evaluation epoch over the validation set.
 
@@ -653,6 +659,12 @@ def validate(
             datamodule: A :class:`~lightning.pytorch.core.datamodule.LightningDataModule` that defines
                 the :class:`~lightning.pytorch.core.hooks.DataHooks.val_dataloader` hook.
 
+            weights_only: Defaults to ``None``. If ``True``, restricts loading to ``state_dicts`` of plain
+                ``torch.Tensor`` and other primitive types. If loading a checkpoint from a trusted source that contains
+                an ``nn.Module``, use ``weights_only=False``. If loading checkpoint from an untrusted source, we
+                recommend using ``weights_only=True``. For more information, please refer to the
+                `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
+
         For more information about multiple dataloaders, see this :ref:`section <multiple-dataloaders>`.
 
         Returns:
@@ -686,17 +698,17 @@ def validate(
         self.state.status = TrainerStatus.RUNNING
         self.validating = True
         return call._call_and_handle_interrupt(
-            self, self._validate_impl, model, dataloaders, ckpt_path, weights_only, verbose, datamodule
+            self, self._validate_impl, model, dataloaders, ckpt_path, verbose, datamodule, weights_only
         )
 
     def _validate_impl(
         self,
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
+        weights_only: Optional[bool] = None,
     ) -> Optional[Union[_PREDICT_OUTPUT, _EVALUATE_OUTPUT]]:
         # --------------------
         # SETUP HOOK
@@ -742,9 +754,9 @@ def test(
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
+        weights_only: Optional[bool] = None,
     ) -> _EVALUATE_OUTPUT:
         r"""Perform one evaluation epoch over the test set. It's separated from fit to make sure you never run on your
         test set until you want to.
@@ -766,6 +778,12 @@ def test(
             datamodule: A :class:`~lightning.pytorch.core.datamodule.LightningDataModule` that defines
                 the :class:`~lightning.pytorch.core.hooks.DataHooks.test_dataloader` hook.
 
+            weights_only: Defaults to ``None``. If ``True``, restricts loading to ``state_dicts`` of plain
+                ``torch.Tensor`` and other primitive types. If loading a checkpoint from a trusted source that contains
+                an ``nn.Module``, use ``weights_only=False``. If loading checkpoint from an untrusted source, we
+                recommend using ``weights_only=True``. For more information, please refer to the
+                `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
+
         For more information about multiple dataloaders, see this :ref:`section <multiple-dataloaders>`.
 
         Returns:
@@ -799,17 +817,17 @@ def test(
         self.state.status = TrainerStatus.RUNNING
         self.testing = True
         return call._call_and_handle_interrupt(
-            self, self._test_impl, model, dataloaders, ckpt_path, weights_only, verbose, datamodule
+            self, self._test_impl, model, dataloaders, ckpt_path, verbose, datamodule, weights_only
         )
 
     def _test_impl(
         self,
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
+        weights_only: Optional[bool] = None,
     ) -> Optional[Union[_PREDICT_OUTPUT, _EVALUATE_OUTPUT]]:
         # --------------------
         # SETUP HOOK
@@ -880,6 +898,12 @@ def predict(
                 Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
                 if a checkpoint callback is configured.
 
+            weights_only: Defaults to ``None``. If ``True``, restricts loading to ``state_dicts`` of plain
+                ``torch.Tensor`` and other primitive types. If loading a checkpoint from a trusted source that contains
+                an ``nn.Module``, use ``weights_only=False``. If loading checkpoint from an untrusted source, we
+                recommend using ``weights_only=True``. For more information, please refer to the
+                `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
+
         For more information about multiple dataloaders, see this :ref:`section <multiple-dataloaders>`.
 
         Returns:
@@ -920,7 +944,7 @@ def predict(
             datamodule,
             return_predictions,
             ckpt_path,
-            weights_only=weights_only,
+            weights_only,
         )
 
     def _predict_impl(