code: def testExtractTensor(self):
      init_dist()
      rank = dist.get_rank()
      device = torch.device(f"cuda:{rank}")
      torch.cuda.set_device(device)
  
      class Model(nn.Module):
          def __init__(self, configs):
              super().__init__()
              self.emb = EmbeddingCollection(
                  tables=configs, device=torch.device("meta")
              )
              # this is incorrect, but doesn't matter.
              self.dense = nn.Linear(10, 10)
  
          def forward(self, x):
              return x
  
      configs = [
          EmbeddingConfig(
              name="AB",
              num_embeddings=4,
              embedding_dim=32,
              feature_names=["A", "B"],
          ),
          EmbeddingConfig(name="C", num_embeddings=8, embedding_dim=32),
      ]
      model = Model(configs=configs)
      model = DMP(module=model, device=device)
  
      plan = model.plan
      sharded_modules = _get_sharded_modules_recursive(
          model.module, "", plan
      )
      sharded_module, params_plan = sharded_modules["emb"]
      ps_collection = PSCollection.fromModule(
          "emb",
          sharded_module,
          params_plan,
          "memory://",
      )
  
      table_names = ps_collection.table_names()
      self.assertEqual(len(table_names), 2)
      self.assertTrue("AB" in table_names)
      self.assertTrue("C" in table_names)
  
file_path: torchrec/contrib/dynamic_embedding/tests/test_ps_collection.py
imports: import unittest
  import torch
  import torch.distributed
  import torch.nn
  from torchrec import EmbeddingCollection, EmbeddingConfig, KeyedJaggedTensor
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec_dynamic_embedding.id_transformer_collection import IDTransformerCollection
  from torchrec_dynamic_embedding.ps import PSCollection
  from torchrec_dynamic_embedding.utils import _get_sharded_modules_recursive
  from utils import init_dist, register_memory_io
label_desc: Use image node_id 1 for calling the TestPSCollection obj's underlying member method code with example usage: obj.testExtractTensor() and returns: x