code: def test_load_state_dict(self) -> None:
      # Set up example KeyedOptimizer 1.
      param_1_t = torch.tensor([1.0, 2.0])
      param_1 = Variable(param_1_t)
      keyed_optimizer_1 = KeyedOptimizer(
          {"param_1": param_1},
          {param_1: {"one": 1.0}},
          [{"params": [param_1], "param_group_val_0": 2.0}],
      )
  
      # Set up example KeyedOptimizer 2.
      param_2_t = torch.tensor([-1.0, -2.0])
      param_2 = Variable(param_2_t)
      keyed_optimizer_2 = KeyedOptimizer(
          {"param_2": param_2},
          {param_2: {"two": -1.0}},
          [{"params": [param_2], "param_group_val_0": -2.0}],
      )
  
      combined_optimizer = CombinedOptimizer(
          [("ko1", keyed_optimizer_1), ("", keyed_optimizer_2)]
      )
      combined_optimizer.save_param_groups(True)
  
      combined_optimizer_state_dict = combined_optimizer.state_dict()
      combined_optimizer_state_dict["state"]["ko1.param_1"] = {
          "one": 999
      }
      combined_optimizer_state_dict["state"]["param_2"] = {"two": 998}
      combined_optimizer_state_dict["param_groups"][0][
          "param_group_val_0"
      ] = 997
      combined_optimizer_state_dict["param_groups"][1][
          "param_group_val_0"
      ] = 996
  
      combined_optimizer.load_state_dict(combined_optimizer_state_dict)
  
      # Check that optimizers in the combined optimizer have their state and
      # param_groups updated.
      self.assertEqual(keyed_optimizer_1.state[param_1], {"one": 999})
      self.assertEqual(keyed_optimizer_2.state[param_2], {"two": 998})
      # pyre-ignore[16]
      self.assertEqual(
          keyed_optimizer_1.param_groups[0]["param_group_val_0"], 997
      )
      self.assertEqual(
          keyed_optimizer_2.param_groups[0]["param_group_val_0"], 996
      )
  
imports: import io
  import os
  import unittest
  from typing import Any, Dict, List
  import torch
  import torch.distributed
  from torch.autograd import Variable
  from torch.distributed._shard import sharded_tensor, sharding_spec
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer, KeyedOptimizerWrapper, OptimizerWrapper
  from torchrec.test_utils import get_free_port
label_desc: Use image node_id 2 for calling the TestCombinedOptimizer obj's underlying member method code with example usage: obj.test_load_state_dict() without return types