Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update src/lightning/fabric/strategies/xla_fsdp.py

This commit is contained in:
Luca Antiga 2024-11-25 11:35:16 +01:00 committed by GitHub
parent 5d227ff480
commit f94efa7108
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/lightning/fabric/strategies/xla_fsdp.py
View File

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