Add @override for files in src/lightning/fabric/plugins/io (#19157)

Author: VictorPrins

src/lightning/fabric/plugins/io/torch_io.py

@@ -15,6 +15,8 @@
 import os
 from typing import Any, Callable, Dict, Optional
 
+from typing_extensions import override
+
 from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
 from lightning.fabric.utilities.cloud_io import _atomic_save, get_filesystem
 from lightning.fabric.utilities.cloud_io import _load as pl_load
@@ -31,6 +33,7 @@ class TorchCheckpointIO(CheckpointIO):
 
     """
 
+    @override
     def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
         """Save model/training states as a checkpoint file through state-dump and file-write.
 
@@ -54,6 +57,7 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_optio
         fs.makedirs(os.path.dirname(path), exist_ok=True)
         _atomic_save(checkpoint, path)
 
+    @override
     def load_checkpoint(
         self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage
     ) -> Dict[str, Any]:
@@ -78,6 +82,7 @@ def load_checkpoint(
 
         return pl_load(path, map_location=map_location)
 
+    @override
     def remove_checkpoint(self, path: _PATH) -> None:
         """Remove checkpoint file from the filesystem.
 

src/lightning/fabric/plugins/io/xla.py

@@ -18,6 +18,7 @@
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
 from lightning_utilities.core.imports import RequirementCache
+from typing_extensions import override
 
 from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
 from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
@@ -39,6 +40,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             raise ModuleNotFoundError(str(_XLA_AVAILABLE))
         super().__init__(*args, **kwargs)
 
+    @override
     def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
         """Save model/training states as a checkpoint file through state-dump and file-write.