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

Code before:
from lightning.fabric.utilities.cloud_io import _load as pl_load
from lightning.fabric.utilities.data import AttributeDict
from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
from lightning.pytorch.accelerators import CUDAAccelerator, MPSAccelerator, XLAAccelerator
from lightning.pytorch.utilities.imports import _OMEGACONF_AVAILABLE
from lightning.pytorch.utilities.migration import pl_legacy_patch
from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
from lightning.pytorch.utilities.model_helpers import is_overridden
from lightning.pytorch.utilities.parsing import parse_class_init_keys
from lightning.pytorch.utilities.rank_zero import rank_zero_warn

if TYPE_CHECKING:
    from torch.storage import UntypedStorage

log = logging.getLogger(__name__)
# the older shall be on the top
CHECKPOINT_PAST_HPARAMS_KEYS = ("hparams", "module_arguments")  # used in 0.7.6


def _load_from_checkpoint(
    cls: Union[type["pl.LightningModule"], type["pl.LightningDataModule"]],
    checkpoint_path: Union[_PATH, IO],
    map_location: _MAP_LOCATION_TYPE = None,
    hparams_file: Optional[_PATH] = None,
    strict: Optional[bool] = None,
    weights_only: Optional[bool] = None,
    **kwargs: Any,
) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
    map_location = map_location or _default_map_location

    if weights_only is None:
        log.debug("`weights_only` not specified, defaulting to `True`.")
        weights_only = True

    with pl_legacy_patch():
        checkpoint = pl_load(checkpoint_path, map_location=map_location, weights_only=weights_only)

    # convert legacy checkpoints to the new format
    checkpoint = _pl_migrate_checkpoint(
        checkpoint, checkpoint_path=(checkpoint_path if isinstance(checkpoint_path, (str, Path)) else None)
    )

    if hparams_file is not None:
        extension = str(hparams_file).split(".")[-1]
        if extension.lower() == "csv":
            hparams = load_hparams_from_tags_csv(hparams_file)
        elif extension.lower() in ("yml", "yaml"):
            hparams = load_hparams_from_yaml(hparams_file)
        else:
            raise ValueError(".csv, .yml or .yaml is required for `hparams_file`")


Review comment (refactor):
Should we default to `weights_only=None` or `weights_only=True`? If we have no use for `weights_only=None`, we can simplify the type hint to `weights_only: bool = True`.

Code after:
from lightning.fabric.utilities.cloud_io import _load as pl_load
from lightning.fabric.utilities.data import AttributeDict
from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
from lightning.pytorch.accelerators import CUDAAccelerator, MPSAccelerator, XLAAccelerator
from lightning.pytorch.utilities.imports import _OMEGACONF_AVAILABLE
from lightning.pytorch.utilities.migration import pl_legacy_patch
from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
from lightning.pytorch.utilities.model_helpers import is_overridden
from lightning.pytorch.utilities.parsing import parse_class_init_keys
from lightning.pytorch.utilities.rank_zero import rank_zero_warn

if TYPE_CHECKING:
    from torch.storage import UntypedStorage

log = logging.getLogger(__name__)
# the older shall be on the top
CHECKPOINT_PAST_HPARAMS_KEYS = ("hparams", "module_arguments")  # used in 0.7.6


def _load_from_checkpoint(
    cls: Union[type["pl.LightningModule"], type["pl.LightningDataModule"]],
    checkpoint_path: Union[_PATH, IO],
    map_location: _MAP_LOCATION_TYPE = None,
    hparams_file: Optional[_PATH] = None,
    strict: Optional[bool] = None,
    weights_only: Optional[bool] = None,
    **kwargs: Any,
) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
    map_location = map_location or _default_map_location

    with pl_legacy_patch():
        checkpoint = pl_load(checkpoint_path, map_location=map_location, weights_only=weights_only)

    # convert legacy checkpoints to the new format
    checkpoint = _pl_migrate_checkpoint(
        checkpoint, checkpoint_path=(checkpoint_path if isinstance(checkpoint_path, (str, Path)) else None)
    )

    if hparams_file is not None:
        extension = str(hparams_file).split(".")[-1]
        if extension.lower() == "csv":
            hparams = load_hparams_from_tags_csv(hparams_file)
        elif extension.lower() in ("yml", "yaml"):
            hparams = load_hparams_from_yaml(hparams_file)
        else:
            raise ValueError(".csv, .yml or .yaml is required for `hparams_file`")

        # overwrite hparams by the given file
        checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = hparams

    # TODO: make this a migration: