code: def test_optimizer_state_on_device(tmpdir):
      """Test that optimizers that create state initially at instantiation still end up with the state on the GPU."""
  
      class TestModel(BoringModel):
          def configure_optimizers(self):
              # Adagrad creates state tensors immediately, model is not yet on GPU.
              return optim.Adagrad(self.parameters())
  
          def on_train_start(self, *args, **kwargs):
              opt = self.optimizers()
              _, state = next(iter(opt.state.items()))
              assert (
                  state["sum"].device
                  == torch.device("cuda", self.local_rank)
                  == self.device
              )
  
      model = TestModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          accelerator="gpu",
          devices=2,
          strategy="ddp",
          fast_dev_run=True,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
  
imports: from unittest import mock
  from unittest.mock import call
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.callbacks import ModelCheckpoint
  from lightning.pytorch.core.optimizer import _configure_optimizers, _configure_schedulers_automatic_opt, _init_optimizers_and_lr_schedulers
  from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.types import LRSchedulerConfig
  from torch import optim
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 21 for calling a global function with example usage: test_optimizer_state_on_device(tmpdir) and returns: optim