code: def setup(config):
      data_dir = config.get("data_dir", None)
      seed = config.get("seed", 42)
      batch_size = config.get("batch_size", 64)
      use_adasum = config.get("use_adasum", False)
      lr = config.get("lr", 0.01)
      momentum = config.get("momentum", 0.5)
      use_cuda = config.get("use_cuda", False)
  
      # 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,
      )
  
      return model, optimizer, train_loader, train_sampler
  
imports: import argparse
  import os
  import horovod.torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torch.utils.data.distributed
  from filelock import FileLock
  from torchvision import datasets, transforms
  import ray
  from ray import train
  from ray.train import ScalingConfig
  from ray.train.horovod import HorovodTrainer
label_desc: Use image node_id 2 for calling a global function with example usage: setup(config) and returns: model,  optimizer,  train_loader,  train_sampler