code: def test_optimizer_to_device():
      @dataclasses.dataclass(frozen=True)
      class FooState:
          bar: int
  
      class TestOptimizer(torch.optim.SGD):
          def __init__(self, *args, **kwargs):
              super().__init__(*args, **kwargs)
              self.state["dummy"] = torch.tensor(0)
              self.state["frozen"] = FooState(0)
  
      layer = torch.nn.Linear(32, 2)
      opt = TestOptimizer(layer.parameters(), lr=0.1)
      _optimizer_to_device(opt, "cpu")
      if torch.cuda.is_available():
          _optimizer_to_device(opt, "cuda")
          assert_opt_parameters_on_device(opt, "cuda")
  
imports: import collections
  import dataclasses
  import torch
  from lightning.fabric.utilities.optimizer import _optimizer_to_device
  from torch import Tensor
label_desc: Use image node_id 1 for calling a global function with example usage: test_optimizer_to_device() without return types