code: def train(gpu, args):
      rank = args.nr * args.gpus + gpu
      dist.init_process_group(
          backend="nccl",
          init_method="env://",
          world_size=args.world_size,
          rank=rank,
      )
      torch.manual_seed(0)
      model = ConvNet()
      torch.cuda.set_device(gpu)
      model.cuda(gpu)
      batch_size = 100
      # define loss function (criterion) and optimizer
      criterion = nn.CrossEntropyLoss().cuda(gpu)
      optimizer = torch.optim.SGD(model.parameters(), 1e-4)
      # Wrap the model
  
      model = DDP(model, device_ids=[gpu])
      # Data loading code
      train_dataset = torchvision.datasets.MNIST(
          root="./data",
          train=True,
          transform=transforms.ToTensor(),
          download=True,
      )
      train_sampler = torch.utils.data.distributed.DistributedSampler(
          train_dataset, num_replicas=args.world_size, rank=rank
      )
      train_loader = torch.utils.data.DataLoader(
          dataset=train_dataset,
          batch_size=batch_size,
          shuffle=False,
          num_workers=0,
          pin_memory=True,
          sampler=train_sampler,
      )
  
      start = datetime.now()
      total_step = len(train_loader)
      for epoch in range(args.epochs):
          for i, (images, labels) in enumerate(train_loader):
              images = images.cuda(non_blocking=True)
              labels = labels.cuda(non_blocking=True)
              # Forward pass
              outputs = model(images)
              loss = criterion(outputs, labels)
  
              # Backward and optimize
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              if (i + 1) % 100 == 0 and gpu == 0:
                  print(
                      "Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}".format(
                          epoch + 1,
                          args.epochs,
                          i + 1,
                          total_step,
                          loss.item(),
                      )
                  )
      if gpu == 0:
          print("Training complete in: " + str(datetime.now() - start))
  
file_path: byteps/example/pytorch/mnist-distributed.py
imports: import os
  from datetime import datetime
  import argparse
  import torch.multiprocessing
  import torchvision
  import torchvision.transforms
  import torch
  import torch.nn
  import torch.distributed
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 2 for calling a global function with example usage: train(gpu, args) without return types