load_from_checkpoint returns the expected type

src/pytorch_lightning/core/datamodule.py

@@ -17,6 +17,7 @@
 from typing import Any, Dict, IO, List, Mapping, Optional, Sequence, Tuple, Union
 
 from torch.utils.data import DataLoader, Dataset, IterableDataset
+from typing_extensions import Self
 
 import pytorch_lightning as pl
 from lightning_lite.utilities.types import _PATH
@@ -218,7 +219,7 @@ def load_from_checkpoint(
         checkpoint_path: Union[_PATH, IO],
         hparams_file: Optional[_PATH] = None,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Self:  # type: ignore [valid-type]
         r"""
         Primary way of loading a datamodule from a checkpoint. When Lightning saves a checkpoint
         it stores the arguments passed to ``__init__``  in the checkpoint under ``"datamodule_hyper_parameters"``.

src/pytorch_lightning/core/saving.py

@@ -26,6 +26,7 @@
 
 import yaml
 from lightning_utilities.core.apply_func import apply_to_collection
+from typing_extensions import Self
 
 import pytorch_lightning as pl
 from lightning_lite.utilities.cloud_io import _load as pl_load
@@ -63,7 +64,7 @@ def load_from_checkpoint(
         hparams_file: Optional[str] = None,
         strict: bool = True,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Self:  # type: ignore [valid-type]
         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"``.