Update fsdp.py

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: