code: def _assert_state_dict_equals(
      self, dict1: Dict[str, Any], dict2: Dict[str, Any]
  ) -> None:
      self.assertEqual(dict1["param_groups"], dict2["param_groups"])
      self.assertEqual(
          dict1["state"]["param_2"],
          dict2["state"]["param_2"],
      )
      torch.testing.assert_close(
          dict1["state"]["param_1"]["tensor"],
          dict2["state"]["param_1"]["tensor"],
      )
      torch.testing.assert_close(
          dict1["state"]["param_1"]["nested_dictionary"]["tensor"],
          dict2["state"]["param_1"]["nested_dictionary"]["tensor"],
      )
      torch.testing.assert_close(
          dict1["state"]["param_1"]["optimizer_module"]["tensor"],
          dict2["state"]["param_1"]["optimizer_module"]["tensor"],
      )
  
      torch.testing.assert_close(
          dict1["state"]["param_1"]["sharded_tensor"]
          .local_shards()[0]
          .tensor,
          dict2["state"]["param_1"]["sharded_tensor"]
          .local_shards()[0]
          .tensor,
      )
  
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 TestKeyedOptimizer obj's underlying member method code with example usage: obj._assert_state_dict_equals(dict1, dict2) without return types