Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/fabric.py (Python)

Code before:
            if not isinstance(filter, dict):
                raise TypeError(f"Filter should be a dictionary, given {filter!r}")
            if not set(filter).issubset(state):
                raise ValueError(
                    f"The filter keys {filter.keys() - state} are not present in the state keys {set(state)}."
                )
            for k, v in filter.items():
                if not callable(v):
                    raise TypeError(f"Expected `fabric.save(filter=...)` for key {k!r} to be a callable, given {v!r}")
        self._strategy.save_checkpoint(path=path, state=_unwrap_objects(state), filter=filter)
        self.barrier()

    def load(
        self,
        path: Union[str, Path],
        state: Optional[dict[str, Union[nn.Module, Optimizer, Any]]] = None,
        strict: bool = True,
        *,
        weights_only: Optional[bool] = None,
    ) -> dict[str, Any]:
        """Load a checkpoint from a file and restore the state of objects (modules, optimizers, etc.) How and which
        processes load gets determined by the `strategy`.

        This method must be called on all processes!

        """
        unwrapped_state = _unwrap_objects(state)
        remainder = self._strategy.load_checkpoint(
            path=path,
            state=unwrapped_state,
            strict=strict,
            weights_only=weights_only,
        )
        self.barrier()
        if state is not None:
            for k in list(unwrapped_state.keys()):
                obj, _ = _unwrap_compiled(state[k])
                if isinstance(obj, (_FabricModule, _FabricOptimizer, _FabricDataLoader)):
                    continue
                state[k] = unwrapped_state[k]
        return remainder

    def load_raw(self, path: Union[str, Path], obj: Union[nn.Module, Optimizer], strict: bool = True) -> None:
        """Load the state of a module or optimizer from a single state-dict file.

        Use this for loading a raw PyTorch model checkpoint created without Fabric.
        This is conceptually equivalent to ``obj.load_state_dict(torch.load(path))``, but is agnostic to the strategy
        being used.

        Args:
            path: A path to where the file is located

Review comment (suggestion):
The original docstring should be extended with the additional kwarg `weight_only` instead of removed.

Code after:
            if not isinstance(filter, dict):
                raise TypeError(f"Filter should be a dictionary, given {filter!r}")
            if not set(filter).issubset(state):
                raise ValueError(
                    f"The filter keys {filter.keys() - state} are not present in the state keys {set(state)}."
                )
            for k, v in filter.items():
                if not callable(v):
                    raise TypeError(f"Expected `fabric.save(filter=...)` for key {k!r} to be a callable, given {v!r}")
        self._strategy.save_checkpoint(path=path, state=_unwrap_objects(state), filter=filter)
        self.barrier()

    def load(
        self,
        path: Union[str, Path],
        state: Optional[dict[str, Union[nn.Module, Optimizer, Any]]] = None,
        strict: bool = True,
        *,
        weights_only: Optional[bool] = None,
    ) -> dict[str, Any]:
        """Load a checkpoint from a file and restore the state of objects (modules, optimizers, etc.).

        How and which processes load gets determined by the `strategy`.
        This method must be called on all processes!

        Args:
            path: A path to where the file is located.
            state: A dictionary of objects whose state will be restored in-place from the checkpoint path.
                If no state is given, then the checkpoint will be returned in full.
            strict: Whether to enforce that the keys in ``state`` match the keys in the checkpoint.
            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>`_.

        Returns:
            The remaining items that were not restored into the given state dictionary. If no state dictionary is
            given, the full checkpoint will be returned.

        Example::

            # Load full checkpoint
            checkpoint = fabric.load("checkpoint.pth")

            # Load into existing objects
            state = {"model": model, "optimizer": optimizer}
            remainder = fabric.load("checkpoint.pth", state)
            epoch = remainder.get("epoch", 0)

        """