code: def main(args):
      # get PyTorch environment variables
      world_size = int(os.environ["WORLD_SIZE"])
      rank = int(os.environ["RANK"])
      local_rank = int(os.environ["LOCAL_RANK"])
  
      distributed = world_size > 1
  
      # set device
      if distributed:
          device = torch.device("cuda", local_rank)
      else:
          device = torch.device(
              "cuda:0" if torch.cuda.is_available() else "cpu"
          )
  
      # initialize distributed process group using default env:// method
      if distributed:
          torch.distributed.init_process_group(backend="nccl")
  
      # define train and test dataset DataLoaders
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
          ]
      )
  
      train_set = torchvision.datasets.CIFAR10(
          root=args.data_dir,
          train=True,
          download=False,
          transform=transform,
      )
  
      if distributed:
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(train_set)
          )
      else:
          train_sampler = None
  
      train_loader = torch.utils.data.DataLoader(
          train_set,
          batch_size=args.batch_size,
          shuffle=(train_sampler is None),
          num_workers=args.workers,
          sampler=train_sampler,
      )
  
      test_set = torchvision.datasets.CIFAR10(
          root=args.data_dir,
          train=False,
          download=False,
          transform=transform,
      )
      test_loader = torch.utils.data.DataLoader(
          test_set,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.workers,
      )
  
      model = Net().to(device)
  
      # wrap model with DDP
      if distributed:
          model = nn.parallel.DistributedDataParallel(
              model, device_ids=[local_rank], output_device=local_rank
          )
  
      # define loss function and optimizer
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(
          model.parameters(),
          lr=args.learning_rate,
          momentum=args.momentum,
      )
  
      # train the model
      for epoch in range(args.epochs):
          print("Rank %d: Starting epoch %d" % (rank, epoch))
          if distributed:
              train_sampler.set_epoch(epoch)
          model.train()
          train(
              train_loader,
              model,
              criterion,
              optimizer,
              epoch,
              device,
              args.print_freq,
              rank,
          )
  
      print("Rank %d: Finished Training" % (rank))
  
      if not distributed or rank == 0:
          os.makedirs(args.output_dir, exist_ok=True)
          model_path = os.path.join(args.output_dir, "cifar_net.pt")
          torch.save(model.state_dict(), model_path)
  
          # evaluate on full test dataset
          evaluate(test_loader, model, device)
  
file_path: MachineLearningNotebooks/how-to-use-azureml/ml-frameworks/pytorch/distributed-pytorch-with-distributeddataparallel/train.py
imports: import torch
  import torchvision
  import torchvision.transforms
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import os
  import argparse
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types