code: def test_opt_state_correct(self) -> None:
      num_features = 4
  
      tables = [
          EmbeddingBagConfig(
              num_embeddings=(i + 1) * 10,
              embedding_dim=(i + 1) * 4,
              name="table_" + str(i),
              feature_names=["feature_" + str(i)],
          )
          for i in range(num_features)
      ]
      ebc = EmbeddingBagCollection(
          tables=tables, device=torch.device("meta")
      )
      apply_optimizer_in_backward(
          RowWiseAdagrad,
          ebc.parameters(),
          {"lr": 0.01},
      )
      plan = construct_module_sharding_plan(
          ebc,
          apply_to_all(ebc, table_wise(rank=0)),
      )
      ebc = shard(
          module=ebc,
          plan=plan,
          device=self.curr_device,
      )
      for name, param in ebc.named_parameters():
          table_name = name[len("embedding_bags.") : -len("weight") - 1]
          self.assertEqual(
              param._in_backward_optimizers[0]
              .state_dict()["state"][""][f"{table_name}.momentum1"]
              .local_tensor()
              .data_ptr(),
              ebc._optim.state_dict()["state"][
                  f"embedding_bags.{table_name}.weight"
              ][f"{table_name}.momentum1"]
              .local_tensor()
              .data_ptr(),
          )
  
file_path: torchrec/torchrec/distributed/composable/tests/test_fused_optim.py
imports: import os
  import unittest
  import torch
  from torch import distributed
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec.distributed.shard import shard
  from torchrec.distributed.sharding_plan import apply_to_all, construct_module_sharding_plan, table_wise
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
  from torchrec.optim.warmup import WarmupOptimizer, WarmupPolicy, WarmupStage
  from torchrec.test_utils import get_free_port
label_desc: Use image node_id 3 for calling the TestFusedOptim obj's underlying member method code with example usage: obj.test_opt_state_correct() without return types