code: def main(argv: List[str]) -> None:
      # parse command line arguments
      args = get_args(argv)
  
      # seed the RNG for better reproducibility. see docs https://pytorch.org/docs/stable/notes/randomness.html
      seed(args.seed)
  
      # device and process group initialization
      device = init_from_env()
  
      # avoid torch autocast exception
      if device.type == "mps":
          device = torch.device("cpu")
  
      path = tempfile.mkdtemp()
      tb_logger = TensorBoardLogger(path)
  
      on_cuda = device.type == "cuda"
  
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
  
      train_dataset = datasets.MNIST(
          "../data", train=True, download=True, transform=transform
      )
      eval_dataset = datasets.MNIST(
          "../data", train=False, transform=transform
      )
  
      train_dataloader = DataLoader(
          train_dataset, batch_size=args.batch_size, pin_memory=on_cuda
      )
      eval_dataloader = DataLoader(
          eval_dataset,
          batch_size=args.test_batch_size,
          pin_memory=on_cuda,
      )
  
      module = Net()
      train_accuracy = MulticlassAccuracy(device=device)
  
      my_unit = MyUnit(
          tb_logger=tb_logger,
          train_accuracy=train_accuracy,
          log_every_n_steps=args.log_every_n_steps,
          lr=args.lr,
          gamma=args.gamma,
          module=module,
          device=device,
          strategy="ddp",
          precision=args.precision,
          gradient_accumulation_steps=4,
          detect_anomaly=True,
          clip_grad_norm=1.0,
      )
  
      fit(
          my_unit,
          train_dataloader=train_dataloader,
          eval_dataloader=eval_dataloader,
          max_epochs=args.max_epochs,
          max_train_steps_per_epoch=args.max_train_steps_per_epoch,
      )
  
      if args.save_model:
          torch.save(module.state_dict(), "mnist_cnn.pt")
  
imports: import argparse
  import sys
  import tempfile
  from argparse import Namespace
  from typing import List, Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.optim import Adadelta
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.data import DataLoader
  from torcheval.metrics import MulticlassAccuracy
  from torchtnt.framework.auto_unit import AutoUnit, TrainStepResults
  from torchtnt.framework.fit import fit
  from torchtnt.framework.state import State
  from torchtnt.utils import init_from_env, seed, TLRScheduler
  from torchtnt.utils.loggers import TensorBoardLogger
  from torchvision import datasets, transforms
label_desc: Use image node_id 1 for calling a global function with example usage: main(argv) without return types