code: def test_sharded_fused_optimizer_state_dict(self) -> None:
      WORLD_SIZE = 2
      LOCAL_SIZE = 2
      tables = [
          EmbeddingBagConfig(
              num_embeddings=64,
              embedding_dim=64,
              name="table_" + str(i),
              feature_names=["feature_" + str(i)],
              # test different optimizer table datatypes to ensure optimizer dtype is consistent
              data_type=DataType.FP16 if i > 1 else DataType.FP32,
          )
          for i in range(4)
      ]
  
      self._run_multi_process_test(
          callable=self._test_sharded_fused_optimizer_state_dict,
          tables=tables,
          backend="nccl",
          local_size=LOCAL_SIZE,
          world_size=WORLD_SIZE,
      )
  
file_path: torchrec/torchrec/distributed/composable/tests/test_fused_optim_nccl.py
imports: import unittest
  from typing import List
  import torch
  from torchrec.distributed.shard import shard
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, row_wise
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext, MultiProcessTestBase
  from torchrec.modules.embedding_configs import DataType, EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
  from torchrec.optim.optimizers import PartialRowWiseAdam
  from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
label_desc: Use image node_id 2 for calling the ShardedFusedOptimizerStateDictTest obj's underlying member method code with example usage: obj.test_sharded_fused_optimizer_state_dict() without return types