Update fsdp.py

Author: amorehead

src/lightning/fabric/strategies/fsdp.py

@@ -269,7 +269,7 @@ def setup_environment(self) -> None:
     @override
     def setup_module_and_optimizers(
         self, module: Module, optimizers: List[Optimizer], scheduler: Optional[_LRScheduler] = None
-    ) -> Tuple[Module, List[Optimizer]]:
+    ) -> Tuple[Module, List[Optimizer], Optional[_LRScheduler]]:
         """Wraps the model into a :class:`~torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel`
         module and sets `use_orig_params=True` to keep the reference to the original parameters in the optimizer."""
         use_orig_params = self._fsdp_kwargs.get("use_orig_params")
@@ -281,7 +281,7 @@ def setup_module_and_optimizers(
                 " call `setup_optimizer`."
             )
         module = self.setup_module(module)
-        return module, optimizers
+        return module, optimizers, scheduler
 
     @override
     def setup_module(self, module: Module) -> Module: