code: def test_load_state_dict(self) -> None:
      param_1_t = torch.tensor([1.0, 2.0])
      param_1 = Variable(param_1_t)
      keyed_optimizer = KeyedOptimizer(
          {"param_1": param_1},
          {param_1: {"one": 1.0}},
          [{"params": [param_1], "param_group_val_0": 2.0}],
      )
      optimizer_wrapper = OptimizerWrapper(keyed_optimizer)
      optimizer_wrapper.save_param_groups(True)
  
      optimizer_wrapper_state_dict = optimizer_wrapper.state_dict()
      optimizer_wrapper_state_dict["state"]["param_1"] = {"one": 999}
      optimizer_wrapper_state_dict["param_groups"][0][
          "param_group_val_0"
      ] = 998
      optimizer_wrapper.load_state_dict(optimizer_wrapper_state_dict)
  
      # Check that both keyed_optimizer and optimizer_wrapper have their state and
      # param_groups updated.
      self.assertEqual(keyed_optimizer.state[param_1], {"one": 999})
      self.assertEqual(optimizer_wrapper.state[param_1], {"one": 999})
      # pyre-ignore[16]
      self.assertEqual(
          keyed_optimizer.param_groups[0]["param_group_val_0"], 998
      )
      self.assertEqual(
          optimizer_wrapper.param_groups[0]["param_group_val_0"], 998
      )
  
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 1 for calling the TestOptimizerWrapper obj's underlying member method code with example usage: obj.test_load_state_dict() without return types