code: def setUp(self) -> None:
      os.environ["RANK"] = "0"
      os.environ["WORLD_SIZE"] = "1"
      os.environ["LOCAL_WORLD_SIZE"] = "1"
      os.environ["MASTER_ADDR"] = str("localhost")
      os.environ["MASTER_PORT"] = str(get_free_port())
      os.environ["NCCL_SOCKET_IFNAME"] = "lo"
      if torch.cuda.is_available():
          self.curr_device = torch.device("cuda:0")
          torch.cuda.set_device(self.curr_device)
          backend = "nccl"
      else:
          self.curr_device = torch.device("cpu")
          backend = "gloo"
      dist.init_process_group(backend=backend)
  
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 1 for calling the TestFusedOptim obj's underlying member method code with example usage: obj.setUp() without return types