datamodule weights_only args

Author: matsumotosan

src/lightning/pytorch/core/datamodule.py

@@ -177,6 +177,7 @@ def load_from_checkpoint(
         checkpoint_path: Union[_PATH, IO],
         map_location: _MAP_LOCATION_TYPE = None,
         hparams_file: Optional[_PATH] = None,
+        weights_only: Optional[bool] = None,
         **kwargs: Any,
     ) -> Self:
         r"""Primary way of loading a datamodule from a checkpoint. When Lightning saves a checkpoint it stores the
@@ -206,6 +207,11 @@ def load_from_checkpoint(
                 If your datamodule's ``hparams`` argument is :class:`~argparse.Namespace`
                 and ``.yaml`` file has hierarchical structure, you need to refactor your datamodule to treat
                 ``hparams`` as :class:`~dict`.
+            weights_only: 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>`_.
             \**kwargs: Any extra keyword args needed to init the datamodule. Can also be used to override saved
                 hyperparameter values.
 
@@ -242,6 +248,7 @@ def load_from_checkpoint(
             map_location=map_location,
             hparams_file=hparams_file,
             strict=None,
+            weights_only=weights_only,
             **kwargs,
         )
         return cast(Self, loaded)

src/lightning/pytorch/core/module.py

@@ -1724,6 +1724,11 @@ def load_from_checkpoint(
             strict: Whether to strictly enforce that the keys in :attr:`checkpoint_path` match the keys
                 returned by this module's state dict. Defaults to ``True`` unless ``LightningModule.strict_loading`` is
                 set, in which case it defaults to the value of ``LightningModule.strict_loading``.
+            weights_only: 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>`_.
             \**kwargs: Any extra keyword args needed to init the model. Can also be used to override saved
                 hyperparameter values.