code: def setUp(self) -> None:
      compute_device = "cuda"
      self.topology = Topology(
          world_size=2, compute_device=compute_device
      )
      tables = [
          EmbeddingBagConfig(
              num_embeddings=100 + i,
              embedding_dim=4 * (10 + i),
              name="table_" + str(i),
              feature_names=["feature_" + str(i)],
          )
          for i in range(1)
      ]
      weighted_tables = [
          EmbeddingBagConfig(
              num_embeddings=200 + i,
              embedding_dim=8 * (10 + i),
              name="weighted_table_" + str(i),
              feature_names=["weighted_feature_" + str(i)],
          )
          for i in range(3)
      ]
      self.topology = Topology(
          world_size=2,
          compute_device=compute_device,
          hbm_cap=2000 * 1024**2,
      )
      self.model = TestSparseNN(
          tables=tables, weighted_tables=weighted_tables
      )
      self.enumerator = EmbeddingEnumerator(
          topology=self.topology, batch_size=BATCH_SIZE
      )
  
      self.sharding_options = self.enumerator.enumerate(
          module=self.model, sharders=[TWSharder()]
      )
      for sharding_option in self.sharding_options:
          sharding_option.shards[0].perf = Perf(
              fwd_compute=40, fwd_comms=30, bwd_compute=20, bwd_comms=10
          )
          sharding_option.shards[0].storage = Storage(
              hbm=10 * 1024**2, ddr=1000 * 1024**2
          )
  
imports: import copy
  import unittest
  from typing import cast, List
  from unittest.mock import MagicMock
  from torch import nn
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.planner.constants import BATCH_SIZE
  from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
  from torchrec.distributed.planner.partitioners import GreedyPerfPartitioner, MemoryBalancedPartitioner, OrderedDeviceHardware
  from torchrec.distributed.planner.types import DeviceHardware, ParameterConstraints, PartitionByType, Perf, PlannerError, Shard, ShardingOption, Storage, Topology
  from torchrec.distributed.planner.utils import reset_shard_rank
  from torchrec.distributed.test_utils.test_model import TestSparseNN
  from torchrec.distributed.types import ModuleSharder, ShardingType
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
label_desc: Use image node_id 1 for calling the TestBalanceModules obj's underlying member method code with example usage: obj.setUp() without return types