code: def train_fn():
      # Horovod: initialize library.
      hvd.init()
      torch.manual_seed(args.seed)
      args.cuda = not args.no_cuda and torch.cuda.is_available()
  
      if args.cuda:
          # Horovod: pin GPU to local rank.
          torch.cuda.set_device(hvd.local_rank())
          torch.cuda.manual_seed(args.seed)
  
      # Horovod: limit # of CPU threads to be used per worker.
      torch.set_num_threads(1)
  
      kwargs = (
          {"num_workers": 1, "pin_memory": True} if args.cuda else {}
      )
      data_dir = args.data_dir or "./data"
      with FileLock(os.path.expanduser("~/.horovod_lock")):
          train_dataset = datasets.MNIST(
              data_dir,
              train=True,
              download=True,
              transform=transforms.Compose(
                  [
                      transforms.ToTensor(),
                      transforms.Normalize((0.1307,), (0.3081,)),
                  ]
              ),
          )
      # Horovod: use DistributedSampler to partition the training data.
      train_sampler = torch.utils.data.distributed.DistributedSampler(
          train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
      )
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=args.batch_size,
          sampler=train_sampler,
          **kwargs,
      )
      transformations = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
  
      test_dataset = datasets.MNIST(
          data_dir, train=False, transform=transformations
      )
      # Horovod: use DistributedSampler to partition the test data.
      test_sampler = torch.utils.data.distributed.DistributedSampler(
          test_dataset, num_replicas=hvd.size(), rank=hvd.rank()
      )
      test_loader = torch.utils.data.DataLoader(
          test_dataset,
          batch_size=args.test_batch_size,
          sampler=test_sampler,
          **kwargs,
      )
  
      model = Net()
  
      # By default, Adasum doesn't need scaling up learning rate.
      lr_scaler = hvd.size() if not args.use_adasum else 1
  
      if args.cuda:
          # Move model to GPU.
          model.cuda()
          # If using GPU Adasum allreduce, scale learning rate by local_size.
          if args.use_adasum and hvd.nccl_built():
              lr_scaler = hvd.local_size()
  
      # Horovod: scale learning rate by lr_scaler.
      optimizer = optim.SGD(
          model.parameters(),
          lr=args.lr * lr_scaler,
          momentum=args.momentum,
      )
  
      # Horovod: (optional) compression algorithm.
      compression = (
          hvd.Compression.fp16
          if args.fp16_allreduce
          else hvd.Compression.none
      )
  
      @hvd.elastic.run
      def train(state):
          # post synchronization event (worker added, worker removed) init ...
          for state.epoch in range(state.epoch, args.epochs + 1):
              state.model.train()
  
              train_sampler.set_epoch(state.epoch)
              steps_remaining = len(train_loader) - state.batch
  
              for state.batch, (data, target) in enumerate(
                  train_loader
              ):
                  if state.batch >= steps_remaining:
                      break
  
                  if args.cuda:
                      data, target = data.cuda(), target.cuda()
                  state.optimizer.zero_grad()
                  output = state.model(data)
                  loss = F.nll_loss(output, target)
                  loss.backward()
                  state.optimizer.step()
                  if state.batch % args.log_interval == 0:
                      # Horovod: use train_sampler to determine
                      # the number of examples in this worker's partition.
                      print(
                          "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                              state.epoch,
                              state.batch * len(data),
                              len(train_sampler),
                              100.0 * state.batch / len(train_loader),
                              loss.item(),
                          )
                      )
                  if (
                      state.batch + 1
                  ) % args.num_batches_per_commit == 0:
                      state.commit()
              state.batch = 0
  
      def test():
          model.eval()
          test_loss = 0.0
          test_accuracy = 0.0
          for data, target in test_loader:
              if args.cuda:
                  data, target = data.cuda(), target.cuda()
              output = model(data)
              # sum up batch loss
              test_loss += F.nll_loss(
                  output, target, size_average=False
              ).item()
              # get the index of the max log-probability
              pred = output.data.max(1, keepdim=True)[1]
              test_accuracy += (
                  pred.eq(target.data.view_as(pred)).cpu().float().sum()
              )
  
          # Horovod: use test_sampler to determine the number of examples in
          # this worker's partition.
          test_loss /= len(test_sampler)
          test_accuracy /= len(test_sampler)
  
          # Horovod: average metric values across workers.
          test_loss = metric_average(test_loss, "avg_loss")
          test_accuracy = metric_average(test_accuracy, "avg_accuracy")
  
          # Horovod: print output only on first rank.
          if hvd.rank() == 0:
              print(
                  "\nTest set: Average loss: {:.4f}, Accuracy: {:.2f}%\n".format(
                      test_loss, 100.0 * test_accuracy
                  )
              )
  
      # Horovod: wrap optimizer with DistributedOptimizer.
      optimizer = hvd.DistributedOptimizer(
          optimizer,
          named_parameters=model.named_parameters(),
          compression=compression,
          op=hvd.Adasum if args.use_adasum else hvd.Average,
      )
  
      # adjust learning rate on reset
      def on_state_reset():
          for param_group in optimizer.param_groups:
              param_group["lr"] = args.lr * hvd.size()
  
      state = hvd.elastic.TorchState(model, optimizer, epoch=1, batch=0)
      state.register_reset_callbacks([on_state_reset])
      train(state)
      test()
  
file_path: horovod/examples/ray/pytorch_ray_elastic.py
imports: import argparse
  import os
  from filelock import FileLock
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  import torch.utils.data.distributed
  import horovod.torch
label_desc: Use image node_id 2 for calling a global function with example usage: train_fn() without return types