From f94efa710862d769ecaa2cd411a3905183e7f7b0 Mon Sep 17 00:00:00 2001
From: Luca Antiga <luca@lightning.ai>
Date: Mon, 25 Nov 2024 11:35:16 +0100
Subject: [PATCH] Update src/lightning/fabric/strategies/xla_fsdp.py

---
 src/lightning/fabric/strategies/xla_fsdp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/lightning/fabric/strategies/xla_fsdp.py b/src/lightning/fabric/strategies/xla_fsdp.py
index aec490a89e..41c63dd01f 100644
--- a/src/lightning/fabric/strategies/xla_fsdp.py
+++ b/src/lightning/fabric/strategies/xla_fsdp.py
@@ -198,7 +198,7 @@ class XLAFSDPStrategy(ParallelStrategy, _Sharded):
     @override
     def setup_module_and_optimizers(
         self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
-    ) -> Tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
+    ) -> tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
         """Returns NotImplementedError since for XLAFSDP optimizer setup must happen after module setup."""
         raise NotImplementedError(
             f"The `{type(self).__name__}` does not support the joint setup of module and optimizer(s)."