update

Author: deependujha

src/lightning/fabric/utilities/registry.py

@@ -36,7 +36,7 @@ def _load_external_callbacks(group: str) -> list[Any]:
         A list of all callbacks collected from external factories.
 
     """
-    factories = entry_points(group=group) if _PYTHON_GREATER_EQUAL_3_10_0 else entry_points().get(group, {})
+    factories = entry_points(group=group) if _PYTHON_GREATER_EQUAL_3_10_0 else entry_points().get(group, {})  # type: ignore[arg-type]
 
     external_callbacks: list[Any] = []
     for factory in factories:

src/lightning/pytorch/strategies/fsdp2.py

@@ -67,10 +67,10 @@
     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
 
 try:
-    from torch.distributed.checkpoint.stateful import Stateful
+    from torch.distributed.checkpoint.stateful import Stateful as _TorchStateful
 except ImportError:
-    # define a no-op base class for compatibility
-    class Stateful:
+
+    class _TorchStateful:  # type: ignore[no-redef]
         pass
 
 
@@ -541,7 +541,7 @@ def _init_fsdp2_mp_policy(mp_policy: Optional["MixedPrecisionPolicy"]) -> Option
 
 
 # Code taken from: https://docs.pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html#saving
-class AppState(Stateful):
+class AppState(_TorchStateful):
     """This is a useful wrapper for checkpointing the Application State. Since this object is compliant with the
     Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the dcp.save/load APIs.