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)."