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

Code before:
        stack = ExitStack()
        stack.enter_context(self.root_device)
        stack.enter_context(precision_init_ctx)
        return stack

    def module_init_context(self, empty_init: Optional[bool] = None) -> AbstractContextManager:
        """A context manager wrapping the model instantiation.

        Here, the strategy can control how the parameters of the model get created (device, dtype) and or apply other
        patches to the model.

        Args:
            empty_init: Whether to initialize the model with empty weights (uninitialized memory).
                If ``None``, the strategy will decide. Some strategies may not support all options.

        """
        precision_module_ctx = self.precision.module_init_context()
        stack = ExitStack()
        stack.enter_context(self.root_device)
        stack.enter_context(_EmptyInit(enabled=bool(empty_init)))
        stack.enter_context(precision_module_ctx)
        return stack

    def setup_module_and_optimizers(
        self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
    ) -> Tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
        """Set up a model and multiple optimizers together.

        The returned objects are expected to be in the same order they were passed in. The default implementation will
        call :meth:`setup_module` and :meth:`setup_optimizer` on the inputs.

        """
        module = self.setup_module(module)
        optimizers = [self.setup_optimizer(optimizer) for optimizer in optimizers]
        return module, optimizers, scheduler

    def setup_module(self, module: Module) -> Module:
        """Performs setup for the model, e.g., by wrapping it by another class."""
        return module

    def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
        """Performs setup for the optimizer, e.g., by wrapping it by another class."""
        return optimizer

    @abstractmethod
    def module_to_device(self, module: Module) -> None:
        """Moves the model to the correct device."""

    def batch_to_device(self, batch: Any, device: Optional[torch.device] = None) -> Any:
        """Moves the batch to the correct device.


Review comment (performance):
```suggestion
    ) -> tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
```

Code after:
        """Controls how tensors get created (device, dtype)."""
        precision_init_ctx = self.precision.tensor_init_context()
        stack = ExitStack()
        stack.enter_context(self.root_device)
        stack.enter_context(precision_init_ctx)
        return stack

    def module_init_context(self, empty_init: Optional[bool] = None) -> AbstractContextManager:
        """A context manager wrapping the model instantiation.

        Here, the strategy can control how the parameters of the model get created (device, dtype) and or apply other
        patches to the model.

        Args:
            empty_init: Whether to initialize the model with empty weights (uninitialized memory).
                If ``None``, the strategy will decide. Some strategies may not support all options.

        """
        precision_module_ctx = self.precision.module_init_context()
        stack = ExitStack()
        stack.enter_context(self.root_device)
        stack.enter_context(_EmptyInit(enabled=bool(empty_init)))
        stack.enter_context(precision_module_ctx)
        return stack

    def setup_module_and_optimizers(
        self, module: Module, optimizers: list[Optimizer], scheduler: Optional["_LRScheduler"] = None
    ) -> tuple[Module, list[Optimizer], Optional["_LRScheduler"]]:
        """Set up a model and multiple optimizers together.

        The returned objects are expected to be in the same order they were passed in. The default implementation will
        call :meth:`setup_module` and :meth:`setup_optimizer` on the inputs.

        """
        module = self.setup_module(module)
        optimizers = [self.setup_optimizer(optimizer) for optimizer in optimizers]
        return module, optimizers, scheduler

    def setup_module(self, module: Module) -> Module:
        """Performs setup for the model, e.g., by wrapping it by another class."""
        return module

    def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
        """Performs setup for the optimizer, e.g., by wrapping it by another class."""
        return optimizer

    @abstractmethod
    def module_to_device(self, module: Module) -> None:
        """Moves the model to the correct device."""

    def batch_to_device(self, batch: Any, device: Optional[torch.device] = None) -> Any: