code: def test_load_state_dict(self) -> None:
      os.environ["MASTER_ADDR"] = str("localhost")
      os.environ["MASTER_PORT"] = str(get_free_port())
      dist.init_process_group("gloo", rank=0, world_size=1)
  
      # Set up example KeyedOptimizer.
      param_1_t, param_2_t = torch.tensor([1.0, 2.0]), torch.tensor(
          [3.0, 4.0]
      )
      param_1, param_2 = Variable(param_1_t), Variable(param_2_t)
      keyed_optimizer = KeyedOptimizer(
          {"param_1": param_1, "param_2": param_2},
          {
              param_1: {
                  "one": 1.0,
                  "tensor": torch.tensor([5.0, 6.0]),
                  "sharded_tensor": sharded_tensor.full(
                      # pyre-ignore [28]
                      sharding_spec.ChunkShardingSpec(
                          dim=0, placements=["rank:0/cpu"]
                      ),
                      (4,),
                      fill_value=1.0,
                  ),
                  "nested_dictionary": {
                      "tensor": torch.tensor([7.0, 8.0]),
                  },
                  "optimizer_module": DummyOptimizerModule(
                      torch.tensor([9.0, 10.0])
                  ),
              },
              param_2: {"two": 2.0},
          },
          [
              {
                  "params": [param_1],
                  "param_group_val_0": 3.0,
                  "param_group_val_1": 4.0,
              },
              {
                  "params": [param_2],
                  "param_group_val_0": 5.0,
                  "param_group_val_1": 6.0,
              },
          ],
      )
      keyed_optimizer.save_param_groups(True)
  
      # Assert state_dict is as expected.
      state: Dict[str, Any] = {
          "param_1": {
              "one": 1.0,
              "tensor": torch.tensor([5.0, 6.0]),
              "sharded_tensor": sharded_tensor.full(
                  # pyre-ignore [28]
                  sharding_spec.ChunkShardingSpec(
                      dim=0, placements=["rank:0/cpu"]
                  ),
                  (4,),
                  fill_value=1.0,
              ),
              "nested_dictionary": {
                  "tensor": torch.tensor([7.0, 8.0]),
              },
              "optimizer_module": {
                  "tensor": torch.tensor([9.0, 10.0]),
              },
          },
          "param_2": {"two": 2.0},
      }
      param_groups: List[Dict[str, Any]] = [
          {
              "params": ["param_1"],
              "param_group_val_0": 3.0,
              "param_group_val_1": 4.0,
          },
          {
              "params": ["param_2"],
              "param_group_val_0": 5.0,
              "param_group_val_1": 6.0,
          },
      ]
      expected_state_dict = {
          "state": state,
          "param_groups": param_groups,
      }
      self._assert_state_dict_equals(
          expected_state_dict, keyed_optimizer.state_dict()
      )
  
      # Modify state dict and call load_state_dict.
      # pyre-ignore [6]
      expected_state_dict["state"]["param_1"]["one"] = 10.0
      # pyre-ignore [6]
      expected_state_dict["state"]["param_1"]["tensor"] = torch.tensor(
          [50.0, 60.0]
      )
      # pyre-ignore [6]
      expected_state_dict["state"]["param_1"][
          "sharded_tensor"
      ] = sharded_tensor.full(
          # pyre-ignore [28]
          sharding_spec.ChunkShardingSpec(
              dim=0, placements=["rank:0/cpu"]
          ),
          (4,),
          fill_value=10.0,
      )
      # pyre-ignore [6]
      expected_state_dict["state"]["param_1"]["nested_dictionary"][
          "tensor"
      ] = torch.tensor([70.0, 80.0])
      # pyre-ignore [6]
      expected_state_dict["state"]["param_1"]["optimizer_module"][
          "tensor"
      ] = torch.tensor([90.0, 100.0])
      # pyre-ignore [6]
      expected_state_dict["param_groups"][0]["param_group_val_0"] = 8.0
      # pyre-ignore [6]
      expected_state_dict["param_groups"][1]["param_group_val_1"] = 9.0
  
      keyed_optimizer.load_state_dict(expected_state_dict)
      self._assert_state_dict_equals(
          expected_state_dict, keyed_optimizer.state_dict()
      )
      dist.destroy_process_group()
  
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 TestKeyedOptimizer obj's underlying member method code with example usage: obj.test_load_state_dict() without return types