code: def assert_opt_parameters_on_device(opt, device: str):
      for param in opt.state.values():
          # Not sure there are any global tensors in the state dict
          if isinstance(param, Tensor):
              assert param.data.device.type == device
          elif isinstance(param, collections.abc.Mapping):
              for subparam in param.values():
                  if isinstance(subparam, Tensor):
                      assert param.data.device.type == device
  
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 2 for calling a global function with example usage: assert_opt_parameters_on_device(opt, device) without return types