load_from_checkpoint returns the expected type

pyproject.toml

@@ -29,6 +29,7 @@ files = [
     "src/pytorch_lightning",
     "src/lightning_lite",
     "src/lightning_app",
+    "test/tests_type_checking",
 ]
 exclude = [
     "src/lightning_app/cli/component-template",

src/pytorch_lightning/core/saving.py

@@ -21,7 +21,7 @@
 from copy import deepcopy
 from enum import Enum
 from pathlib import Path
-from typing import Any, Callable, cast, Dict, IO, MutableMapping, Optional, Type, Union
+from typing import Any, Callable, cast, Dict, IO, MutableMapping, Optional, Type, TypeVar, Union
 from warnings import warn
 
 import yaml
@@ -49,6 +49,9 @@
 # the older shall be on the top
 CHECKPOINT_PAST_HPARAMS_KEYS = ("hparams", "module_arguments")  # used in 0.7.6
 
+LM = TypeVar("LM", bound="pl.LightningModule")
+LDM = TypeVar("LDM", bound="pl.LightningDataModule")
+
 
 class ModelIO:
     CHECKPOINT_HYPER_PARAMS_KEY = "hyper_parameters"
@@ -57,13 +60,13 @@ class ModelIO:
 
     @classmethod
     def load_from_checkpoint(
-        cls,
+        cls: Union[Type["ModelIO"], Type[LM], Type[LDM]],
         checkpoint_path: Union[str, IO],
         map_location: _MAP_LOCATION_TYPE = None,
         hparams_file: Optional[str] = None,
         strict: bool = True,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Union[LM, LDM]:
         r"""
         Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint
         it stores the arguments passed to ``__init__``  in the checkpoint under ``"hyper_parameters"``.
@@ -146,13 +149,13 @@ def load_from_checkpoint(
 
 
 def _load_from_checkpoint(
-    cls: Union[Type["ModelIO"], Type["pl.LightningModule"], Type["pl.LightningDataModule"]],
+    cls: Union[Type["ModelIO"], Type[LM], Type[LDM]],
     checkpoint_path: Union[_PATH, IO],
     map_location: _MAP_LOCATION_TYPE = None,
     hparams_file: Optional[_PATH] = None,
     strict: Optional[bool] = None,
     **kwargs: Any,
-) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+) -> Union[LM, LDM]:
     if map_location is None:
         map_location = cast(_MAP_LOCATION_TYPE, lambda storage, loc: storage)
     with pl_legacy_patch():
@@ -182,9 +185,9 @@ def _load_from_checkpoint(
     checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
 
     if issubclass(cls, pl.LightningDataModule):
-        return _load_state(cls, checkpoint, **kwargs)
+        return cast(LDM, _load_state(cls, checkpoint, **kwargs))
     if issubclass(cls, pl.LightningModule):
-        return _load_state(cls, checkpoint, strict=strict, **kwargs)
+        return cast(LM, _load_state(cls, checkpoint, strict=strict, **kwargs))
     raise NotImplementedError(f"Unsupported {cls}")
 
 

tests/tests_type_checking/test_lightning_module.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.demos.boring_classes import BoringModel
+
+
+def test_load_from_checkpoint_type(tmp_path: Path) -> None:
+    class MyModule(BoringModel):
+        def __init__(self, some_parameter: int):
+            super().__init__()
+            self.save_hyperparameters()
+
+        @property
+        def parameter(self) -> int:
+            return self.hparams.some_parameter
+
+    net = MyModule(some_parameter=42)
+    trainer = Trainer(default_root_dir=str(tmp_path), fast_dev_run=True)
+    trainer.fit(net)
+    checkpoint_path = str(tmp_path / "model.pt")
+    trainer.save_checkpoint(checkpoint_path)
+
+    net_loaded = MyModule.load_from_checkpoint(checkpoint_path)  # type: ignore
+    assert net_loaded.parameter == 42