code: def train_fn(
      data_dir=None,
      seed=42,
      use_cuda=False,
      batch_size=64,
      use_adasum=False,
      lr=0.01,
      momentum=0.5,
      num_epochs=10,
      log_interval=10,
  ):
      # Horovod: initialize library.
      hvd.init()
      torch.manual_seed(seed)
  
      if use_cuda:
          # Horovod: pin GPU to local rank.
          torch.cuda.set_device(hvd.local_rank())
          torch.cuda.manual_seed(seed)
  
      # Horovod: limit # of CPU threads to be used per worker.
      torch.set_num_threads(1)
  
      kwargs = (
          {"num_workers": 1, "pin_memory": True} if use_cuda else {}
      )
      data_dir = 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=batch_size,
          sampler=train_sampler,
          **kwargs,
      )
  
      model = Net()
  
      # By default, Adasum doesn't need scaling up learning rate.
      lr_scaler = hvd.size() if not use_adasum else 1
  
      if use_cuda:
          # Move model to GPU.
          model.cuda()
          # If using GPU Adasum allreduce, scale learning rate by local_size.
          if use_adasum and hvd.nccl_built():
              lr_scaler = hvd.local_size()
  
      # Horovod: scale learning rate by lr_scaler.
      optimizer = optim.SGD(
          model.parameters(), lr=lr * lr_scaler, momentum=momentum
      )
  
      # Horovod: wrap optimizer with DistributedOptimizer.
      optimizer = hvd.DistributedOptimizer(
          optimizer,
          named_parameters=model.named_parameters(),
          op=hvd.Adasum if use_adasum else hvd.Average,
      )
  
      for epoch in range(1, num_epochs + 1):
          model.train()
          # Horovod: set epoch to sampler for shuffling.
          train_sampler.set_epoch(epoch)
          for batch_idx, (data, target) in enumerate(train_loader):
              if use_cuda:
                  data, target = data.cuda(), target.cuda()
              optimizer.zero_grad()
              output = model(data)
              loss = F.nll_loss(output, target)
              loss.backward()
              optimizer.step()
              if batch_idx % 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(
                          epoch,
                          batch_idx * len(data),
                          len(train_sampler),
                          100.0 * batch_idx / len(train_loader),
                          loss.item(),
                      )
                  )
  
file_path: ray/release/ml_user_tests/horovod/horovod_example.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
  from horovod.ray import RayExecutor
label_desc: Use image node_id 2 for calling a global function with example usage: train_fn(data_dir, seed, use_cuda, batch_size, use_adasum, lr, momentum, num_epochs, log_interval) without return types