code: def state_dict_test(optimizer, weight, bias, input):
      def fn_base(optimizer, weight, bias, input):
          optimizer.zero_grad()
          loss = (weight.mv(input) + bias).pow(2).sum()
          loss.backward()
          return loss
  
      fn = functools.partial(fn_base, optimizer, weight, bias, input)
  
      # Prime the optimizer
      for _i in range(5):
          optimizer.step(fn)
      # Clone the weights and construct new optimizer for them
      weight_c = weight.data.clone().requires_grad_()
      bias_c = bias.data.clone().requires_grad_()
      optimizer_c = Adam(
          [weight_c, bias_c], lr=1e-3, precision=optimizer.precision
      )
      fn_c = functools.partial(
          fn_base, optimizer_c, weight_c, bias_c, input
      )
      # Load state dict
      state_dict = deepcopy(optimizer.state_dict())
      optimizer_c.load_state_dict(state_dict)
  
      for group, group_c in zip(
          optimizer.param_groups, optimizer_c.param_groups
      ):
          for p, p_c in zip(group["params"], group_c["params"]):
              assert torch.equal(
                  optimizer.state[p]["exp_avg"],
                  optimizer_c.state[p_c]["exp_avg"],
              )
              assert torch.equal(
                  optimizer.state[p]["exp_avg_sq"],
                  optimizer_c.state[p_c]["exp_avg_sq"],
              )
  
      if optimizer.fp32_param_groups:
          # When using mixed precision, fp32_param_groups are made from FP16 params rather than
          # copied via state_dict, introducing differences between the original optimizer and
          # the copy. Because this test requires that they be the exact same, we copy the
          # fp32 params from the original optimizer to the copy
          optimizer_c.fp32_param_groups = deepcopy(
              optimizer.fp32_param_groups
          )
  
      # Run both optimizations in parallel
      for _i in range(5):
          optimizer.step(fn)
          optimizer_c.step(fn_c)
  
          assert torch.equal(weight, weight_c)
          assert torch.equal(bias, bias_c)
  
imports: from copy import deepcopy
  import functools
  import pytest
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: state_dict_test(optimizer, weight, bias, input) and returns: loss