code: def testExtractTensor(self):
      init_dist()
      rank = dist.get_rank()
      device = torch.device(f"cuda:{rank}")
      torch.cuda.set_device(device)
  
      batch_size = 4
      model1 = Model(
          num_embeddings=1000,
          init_max=1,
          init_min=1,
          batch_size=batch_size,
      )
  
      model2 = Model(
          num_embeddings=100,
          init_max=1,
          init_min=1,
          batch_size=batch_size,
      )
      model2_config = model2.config
  
      model2.dense.weight.data.copy_(model1.dense.weight.data)
      model2.dense.bias.data.copy_(model1.dense.bias.data)
  
      def get_dmp(model):
          topology = Topology(
              world_size=dist.get_world_size(),
              local_world_size=dist.get_world_size(),
              compute_device="cuda",
          )
  
          fused_params = {
              "learning_rate": 1e-1,
              "optimizer": EmbOptimType.ADAM,
              "cache_load_factor": 0.1,
          }
          sharders = [
              EmbeddingBagCollectionSharder(fused_params=fused_params),
              FusedEmbeddingBagCollectionSharder(
                  fused_params=fused_params
              ),
              EmbeddingCollectionSharder(fused_params=fused_params),
          ]
          plan = EmbeddingShardingPlanner(
              topology=topology, constraints=None
          ).plan(model, sharders)
          model = DMP(
              module=model, device=device, plan=plan, sharders=sharders
          )
  
          dense_optimizer = KeyedOptimizerWrapper(
              dict(
                  in_backward_optimizer_filter(model.named_parameters())
              ),
              lambda params: torch.optim.Adam(params, lr=1e-1),
          )
          optimizer = CombinedOptimizer(
              [model.fused_optimizer, dense_optimizer]
          )
  
          return model, optimizer
  
      model1, optimizer1 = get_dmp(model1)
      model2, optimizer2 = get_dmp(model2)
  
      transformer = IDTransformerGroup(
          "memory://",
          model2,
          {"emb": [model2_config]},
          transform_config={"type": "naive"},
          ps_config={"chunk_size": 1024},
      )
  
      def sigmoid_crossentropy(y_true, y_pred):
          ce = nn.BCELoss()(y_pred, y_true)
          return torch.mean(torch.sum(ce, dim=-1))
  
      for i in range(100):
          kjt = KeyedJaggedTensor(
              keys=["id"],
              values=torch.randint(0, 1000, (40,), dtype=torch.long),
              lengths=torch.tensor([10, 10, 10, 10], dtype=torch.long),
          )
          kjts, fetch_handles = transformer.transform({"emb": kjt})
          for handle in fetch_handles:
              handle.wait()
          mapped_kjt = kjts["emb"]
          label = torch.randint(0, 2, (4, 1), device=device).float()
          kjt = kjt.to(device)
          mapped_kjt = mapped_kjt.to(device)
  
          output1 = model1(kjt)
          task_loss1 = sigmoid_crossentropy(label, output1)
  
          task_loss1.backward()
          optimizer1.step()
          optimizer1.zero_grad()
  
          output2 = model2(mapped_kjt)
          task_loss2 = sigmoid_crossentropy(label, output2)
  
          task_loss2.backward()
          optimizer2.step()
          optimizer2.zero_grad()
  
          self.assertTrue(abs((task_loss1 - task_loss2).item()) < 1e-7)
  
imports: import unittest
  import torch
  import torch.distributed
  import torch.nn
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from torchrec import EmbeddingCollection, EmbeddingConfig, KeyedJaggedTensor
  from torchrec.distributed.embedding import EmbeddingCollectionSharder
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.fused_embeddingbag import FusedEmbeddingBagCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
  from torchrec_dynamic_embedding.id_transformer_group import IDTransformerGroup
  from utils import init_dist, register_memory_io
label_desc: Use image node_id 1 for calling the TestPSPrecision obj's underlying member method code with example usage: obj.testExtractTensor() and returns: model,  optimizer, torch