code: def main(argv: List[str]):
      args = parse_args(argv)
      rank = int(os.environ["LOCAL_RANK"])
  
      print("Running with args", args)
  
      if torch.cuda.is_available():
          device: torch.device = torch.device(f"cuda:{rank}")
          backend = "nccl"
          torch.cuda.set_device(device)
      else:
          device: torch.device = torch.device("cpu")
          backend = "gloo"
  
      if not torch.distributed.is_initialized():
          dist.init_process_group(backend=backend)
  
      rank = dist.get_rank()
      world_size = dist.get_world_size()
      num_embeddings_per_feature = None
      if args.num_embeddings_per_feature is not None:
          num_embeddings_per_feature = list(
              map(int, args.num_embeddings_per_feature.split(","))
          )
      train_loader = NvtBinaryDataloader(
          binary_file_path=os.path.join(args.binary_path, "train"),
          batch_size=args.batch_size,
      ).get_dataloader(rank=rank, world_size=world_size)
      val_loader = NvtBinaryDataloader(
          binary_file_path=os.path.join(args.binary_path, "validation"),
          batch_size=args.batch_size,
      ).get_dataloader(rank=rank, world_size=world_size)
  
      test_loader = NvtBinaryDataloader(
          binary_file_path=os.path.join(args.binary_path, "test"),
          batch_size=args.batch_size,
      ).get_dataloader(rank=rank, world_size=world_size)
  
      eb_configs = [
          EmbeddingBagConfig(
              name=f"t_{feature_name}",
              embedding_dim=args.embedding_dim,
              num_embeddings=none_throws(num_embeddings_per_feature)[
                  feature_idx
              ]
              if num_embeddings_per_feature is not None
              else args.num_embeddings,
              feature_names=[feature_name],
          )
          for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)
      ]
  
      train_model = DLRMTrain(
          DLRM(
              embedding_bag_collection=EmbeddingBagCollection(
                  tables=eb_configs, device=torch.device("meta")
              ),
              dense_in_features=len(DEFAULT_INT_NAMES),
              dense_arch_layer_sizes=list(
                  map(int, args.dense_arch_layer_sizes.split(","))
              ),
              over_arch_layer_sizes=list(
                  map(int, args.over_arch_layer_sizes.split(","))
              ),
              dense_device=device,
          ),
      )
  
      train_model = fuse_embedding_optimizer(
          train_model,
          optimizer_type=torchrec.optim.RowWiseAdagrad
          if args.adagrad
          else torch.optim.SGD,
          optimizer_kwargs={"learning_rate": args.learning_rate},
          device=torch.device("meta"),
      )
  
      sharders = cast(
          List[ModuleSharder[nn.Module]],
          [
              FusedEmbeddingBagCollectionSharder(),
          ],
      )
  
      pg = dist.GroupMember.WORLD
      hbm_cap = torch.cuda.get_device_properties(device).total_memory
      local_world_size = trec_dist.comm.get_local_size(world_size)
      model = DistributedModelParallel(
          module=train_model,
          device=device,
          env=trec_dist.ShardingEnv.from_process_group(pg),
          plan=trec_dist.planner.EmbeddingShardingPlanner(
              topology=trec_dist.planner.Topology(
                  world_size=world_size,
                  compute_device=device.type,
                  local_world_size=local_world_size,
                  hbm_cap=hbm_cap,
              ),
              storage_reservation=trec_dist.planner.storage_reservations.HeuristicalStorageReservation(
                  percentage=0.25,
              ),
              batch_size=args.batch_size,
          ).collective_plan(train_model, sharders, pg),
          sharders=sharders,
      )
  
      non_fused_optimizer = KeyedOptimizerWrapper(
          dict(in_backward_optimizer_filter(model.named_parameters())),
          lambda params: torch.optim.Adagrad(
              params, lr=args.learning_rate
          )
          if args.adagrad
          else torch.optim.SGD(params, lr=args.learning_rate),
      )
  
      opt = trec_optim.keyed.CombinedOptimizer(
          [non_fused_optimizer, model.fused_optimizer]
      )
  
      train_pipeline = TrainPipelineSparseDist(
          model,
          opt,
          device,
      )
  
      throughput = ThroughputMetric(
          batch_size=args.batch_size,
          world_size=world_size,
          window_seconds=30,
          warmup_steps=10,
      )
  
      changing_point_steps = (
          TOTAL_TRAINING_SAMPLES
          * args.lr_change_point
          // args.batch_size
          // world_size
      )
      for epoch in range(args.epochs):
          print(f"starting the {epoch} epoch now")
          start_time = time.time()
          it = iter(train_loader)
          step = 0
          losses = []
          while True:
              try:
                  train_pipeline._model.train()
                  loss, _logits, _labels = train_pipeline.progress(it)
  
                  if args.change_lr and step == changing_point_steps:
                      print(
                          f"Changing learning rate to: {args.lr_after_change_point}"
                      )
                      optimizer = train_pipeline._optimizer
                      lr = args.lr_after_change_point
                      for g in optimizer.param_groups:
                          g["lr"] = lr
  
                  throughput.update()
                  losses.append(loss)
  
                  if (
                      step % args.throughput_check_freq_within_epoch
                      == 0
                      and step != 0
                  ):
                      # infra calculation
                      throughput_val = throughput.compute()
                      if rank == 0:
                          print("step", step)
                          print("throughput", throughput_val)
                      losses = []
                  if (
                      step % args.validation_freq_within_epoch == 0
                      and step != 0
                  ):
                      # metrics calculation
                      validation_it = iter(val_loader)
                      auroc_result, accuracy_result, bce_loss = _eval(
                          train_pipeline, validation_it
                      )
                      if rank == 0:
                          print(
                              f"AUROC over validation set: {auroc_result}."
                          )
                          print(
                              f"Accuracy over validation set: {accuracy_result}."
                          )
                          print(
                              "binary cross entropy loss",
                              bce_loss / (args.batch_size),
                          )
  
                      test_it = iter(test_loader)
                      auroc_result, accuracy_result, bce_loss = _eval(
                          train_pipeline, test_it
                      )
                      if rank == 0:
                          print(f"AUROC over test set: {auroc_result}.")
                          print(
                              f"Accuracy over test set: {accuracy_result}."
                          )
  
                  step += 1
  
              except StopIteration:
                  print("Reached stop iteration")
                  break
          train_time = time.time()
          if rank == 0:
              print(
                  f"this epoch training takes {train_time - start_time}"
              )
  
          # eval
          val_it = iter(val_loader)
          auroc_result, accuracy_result, bce_loss = _eval(
              train_pipeline, val_it
          )
          if rank == 0:
              print(f"AUROC over validation set: {auroc_result}.")
              print(f"Accuracy over validation set: {accuracy_result}.")
              print(
                  "binary cross entropy loss over validation set",
                  bce_loss / (args.batch_size),
              )
          # test
          test_it = iter(test_loader)
          auroc_result, accuracy_result, bce_loss = _eval(
              train_pipeline, test_it
          )
          if rank == 0:
              print(f"AUROC over test set: {auroc_result}.")
              print(f"Accuracy over test set: {accuracy_result}.")
              print(
                  "binary cross entropy loss over test set",
                  bce_loss / (args.batch_size),
              )
  
imports: import argparse
  import os
  import sys
  import time
  from typing import cast, Iterator, List, Tuple
  import torch
  import torch.distributed
  import torch.nn
  import torchmetrics
  import torchrec
  import torchrec.distributed
  import torchrec.optim
  from nvt_binary_dataloader import NvtBinaryDataloader
  from pyre_extensions import none_throws
  from torchrec import EmbeddingBagCollection
  from torchrec.datasets.criteo import DEFAULT_CAT_NAMES, DEFAULT_INT_NAMES, TOTAL_TRAINING_SAMPLES
  from torchrec.datasets.utils import Batch
  from torchrec.distributed import TrainPipelineSparseDist
  from torchrec.distributed.fused_embeddingbag import FusedEmbeddingBagCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.types import ModuleSharder
  from torchrec.metrics.throughput import ThroughputMetric
  from torchrec.models.dlrm import DLRM, DLRMTrain
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.fused_embedding_modules import fuse_embedding_optimizer
  from torchrec.optim.keyed import KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
label_desc: Use image node_id 3 for calling a global function with example usage: main(argv) without return types