Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/fabric.py (Python)

Code before:

        if compile_kwargs is not None:
            module = _to_compiled(module, compile_kwargs)
        module = _FabricModule(module, self._strategy, original_module=original_module)

        # Update the _DeviceDtypeModuleMixin's device parameter
        # NOTE: for sharded strategies or manual device placement, there's no single root device
        _update_properties(
            module, device=self.device if move_to_device else next(module.parameters(), torch.tensor(0)).device
        )

        optimizers = [_FabricOptimizer(optimizer, self._strategy, self._callbacks) for optimizer in optimizers]

        self._models_setup += 1

        if hasattr(original_module, "_fabric"):  # this is probably a LightningModule
            original_module._fabric = self
            original_module._fabric_optimizers = optimizers
            if original_module not in self._callbacks:
                self._callbacks.append(original_module)

        self.call("on_after_setup", fabric=self, module=module)

        if optimizers:
            # join both types in a tuple for API convenience
            return (module, *optimizers, scheduler)
        return module

    def setup_module(
        self, module: nn.Module, move_to_device: bool = True, _reapply_compile: bool = True
    ) -> _FabricModule:
        r"""Set up a model for accelerated training or inference.

        This is the same as calling ``.setup(model)`` with no optimizers. It is useful for inference or for certain
        strategies like `FSDP` that require setting up the module before the optimizer can be created and set up.
        See also :meth:`setup_optimizers`.

        Args:
            module: A :class:`torch.nn.Module` to set up
            move_to_device: If set ``True`` (default), moves the model to the correct device. Set this to ``False``
                and alternatively use :meth:`to_device` manually.
            _reapply_compile: If ``True`` (default), and the model was ``torch.compile``d before, the
                corresponding :class:`~torch._dynamo.OptimizedModule` wrapper will be removed and reapplied with the
                same settings after the model was set up by the strategy (e.g., after the model was wrapped by DDP,
                FSDP etc.). Set it to ``False`` if compiling DDP/FSDP is causing issues.
        Returns:
            The wrapped model.

        """
        self._validate_setup_module(module)
        module, compile_kwargs = _unwrap_compiled(module) if _reapply_compile else (module, None)

Review comment (bug):
This is a breaking change, it will cause existing user code to fail, because `scheduler` is returned unconditionally.

Since `scheduler` is `Optional` in the signature, I suggest we only return it if it was not `None` as an argument, so we won't break anyone's code.

Code after:
                module, list(optimizers), scheduler
            )
        else:
            module = self._strategy.setup_module(module)

        if compile_kwargs is not None:
            module = _to_compiled(module, compile_kwargs)
        module = _FabricModule(module, self._strategy, original_module=original_module)

        # Update the _DeviceDtypeModuleMixin's device parameter
        # NOTE: for sharded strategies or manual device placement, there's no single root device
        _update_properties(
            module, device=self.device if move_to_device else next(module.parameters(), torch.tensor(0)).device
        )

        optimizers = [_FabricOptimizer(optimizer, self._strategy, self._callbacks) for optimizer in optimizers]

        self._models_setup += 1

        if hasattr(original_module, "_fabric"):  # this is probably a LightningModule
            original_module._fabric = self
            original_module._fabric_optimizers = optimizers
            if original_module not in self._callbacks:
                self._callbacks.append(original_module)

        self.call("on_after_setup", fabric=self, module=module)

        if optimizers:
            # join both types in a tuple for API convenience
            return (module, *optimizers, scheduler) if scheduler is not None else (module, *optimizers)
        return module

    def setup_module(
        self, module: nn.Module, move_to_device: bool = True, _reapply_compile: bool = True
    ) -> _FabricModule:
        r"""Set up a model for accelerated training or inference.

        This is the same as calling ``.setup(model)`` with no optimizers. It is useful for inference or for certain
        strategies like `FSDP` that require setting up the module before the optimizer can be created and set up.
        See also :meth:`setup_optimizers`.

        Args:
            module: A :class:`torch.nn.Module` to set up
            move_to_device: If set ``True`` (default), moves the model to the correct device. Set this to ``False``
                and alternatively use :meth:`to_device` manually.
            _reapply_compile: If ``True`` (default), and the model was ``torch.compile``d before, the
                corresponding :class:`~torch._dynamo.OptimizedModule` wrapper will be removed and reapplied with the
                same settings after the model was set up by the strategy (e.g., after the model was wrapped by DDP,
                FSDP etc.). Set it to ``False`` if compiling DDP/FSDP is causing issues.
        Returns:
            The wrapped model.