code: def train(args):
      world_size = len(args.hosts)
      is_distributed = world_size > 1
      logger.debug(
          "Number of hosts {}. Distributed training - {}".format(
              world_size, is_distributed
          )
      )
      use_cuda = args.num_gpus > 0
      logger.debug(
          "Number of gpus available - {}".format(args.num_gpus)
      )
      kwargs = (
          {"num_workers": 1, "pin_memory": True} if use_cuda else {}
      )
      device = torch.device("cuda" if use_cuda else "cpu")
  
      if is_distributed:
          # Initialize the distributed environment.
          backend = "gloo"
          os.environ["WORLD_SIZE"] = str(world_size)
          host_rank = args.hosts.index(args.current_host)
          dist.init_process_group(
              backend=backend, rank=host_rank, world_size=world_size
          )
          logger.info(
              "Initialized the distributed environment: '{}' backend on {} nodes. ".format(
                  backend, dist.get_world_size()
              )
              + "Current host rank is {}. Is cuda available: {}. Number of gpus: {}".format(
                  dist.get_rank(),
                  torch.cuda.is_available(),
                  args.num_gpus,
              )
          )
  
      # set the seed for generating random numbers
      seed = 1
      torch.manual_seed(seed)
      if use_cuda:
          torch.cuda.manual_seed(seed)
  
      train_sampler, train_loader = _get_train_data_loader(
          args.data_dir, is_distributed, args.batch_size, **kwargs
      )
      test_loader = _get_test_data_loader(args.data_dir, **kwargs)
  
      logger.debug(
          "Processes {}/{} ({:.0f}%) of train data".format(
              len(train_loader.sampler),
              len(train_loader.dataset),
              100.0
              * len(train_loader.sampler)
              / len(train_loader.dataset),
          )
      )
  
      logger.debug(
          "Processes {}/{} ({:.0f}%) of test data".format(
              len(test_loader.sampler),
              len(test_loader.dataset),
              100.0
              * len(test_loader.sampler)
              / len(test_loader.dataset),
          )
      )
  
      model = Net().to(device)
      if is_distributed and use_cuda:
          # multi-machine multi-gpu case
          logger.debug(
              "Multi-machine multi-gpu: using DistributedDataParallel."
          )
          model = torch.nn.parallel.DistributedDataParallel(model)
      elif use_cuda:
          # single-machine multi-gpu case
          logger.debug(
              "Single-machine multi-gpu: using DataParallel().cuda()."
          )
          model = torch.nn.DataParallel(model)
      else:
          # single-machine or multi-machine cpu case
          logger.debug(
              "Single-machine/multi-machine cpu: using DataParallel."
          )
          model = torch.nn.DataParallel(model)
  
      optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.5)
  
      log_interval = 100
      for epoch in range(1, args.epochs + 1):
          if is_distributed:
              train_sampler.set_epoch(epoch)
          model.train()
          for batch_idx, (data, target) in enumerate(train_loader, 1):
              data, target = data.to(device), target.to(device)
              optimizer.zero_grad()
              output = model(data)
              loss = F.nll_loss(output, target)
              loss.backward()
              if is_distributed and not use_cuda:
                  # average gradients manually for multi-machine cpu case only
                  _average_gradients(model)
              optimizer.step()
              if batch_idx % log_interval == 0:
                  logger.debug(
                      "Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}".format(
                          epoch,
                          batch_idx * len(data),
                          len(train_loader.sampler),
                          100.0 * batch_idx / len(train_loader),
                          loss.item(),
                      )
                  )
          accuracy = test(model, test_loader, device)
      save_model(model, args.model_dir)
  
      logger.debug("Overall test accuracy: {};".format(accuracy))
  
imports: import argparse
  import json
  import logging
  import os
  import sys
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  from torchvision import datasets, transforms
label_desc: Use image node_id 4 for calling a global function with example usage: train(args) without return types