update

Author: deependujha

src/lightning/pytorch/strategies/fsdp2.py

@@ -67,7 +67,13 @@
     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
 
 if _TORCH_GREATER_EQUAL_2_6:
-    from torch.distributed.checkpoint.stateful import Stateful as _TorchStateful
+    try:
+        from torch.distributed.checkpoint.stateful import Stateful as _TorchStateful
+    except ImportError:
+
+        class _TorchStateful:  # type: ignore[no-redef]
+            pass
+
 else:
 
     class _TorchStateful:  # type: ignore[no-redef]