code: def run(large=False):
      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)
  
      # If set > 0, will resume training from a given checkpoint.
      resume_from_epoch = 0
      for try_epoch in range(args.epochs, 0, -1):
          if os.path.exists(
              args.checkpoint_format.format(epoch=try_epoch)
          ):
              resume_from_epoch = try_epoch
              break
  
      # Load MNIST dataset
      train_loader, train_sampler = load_data_mnist()
  
      model = Net(large=large)
      if args.cuda:
          model.cuda()
  
      # Horovod: scale learning rate by the number of GPUs.
      optimizer = optim.SGD(
          model.parameters(),
          lr=args.lr * np.sqrt(hvd.size()),
          momentum=0.9,
          weight_decay=5e-4,
      )
  
      # Horovod: wrap optimizer with DistributedOptimizer.
      optimizer = hvd.DistributedOptimizer(
          optimizer, named_parameters=model.named_parameters()
      )
  
      scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
          optimizer, T_max=200
      )
  
      # Restore from a previous checkpoint, if initial_epoch is specified.
      # Horovod: restore on the first worker which will broadcast
      # weights to other workers.
      if resume_from_epoch > 0 and hvd.rank() == 0:
          filepath = args.checkpoint_format.format(
              epoch=resume_from_epoch
          )
          checkpoint = torch.load(filepath)
          model.load_state_dict(checkpoint["model"])
          optimizer.load_state_dict(checkpoint["optimizer"])
          scheduler.load_state_dict(checkpoint["scheduler"])
  
      def on_state_reset():
          # Horovod: scale the learning rate as controlled by the LR schedule
          scheduler.base_lrs = [
              args.lr * hvd.size() for _ in scheduler.base_lrs
          ]
  
      state = hvd.elastic.TorchState(
          model=model,
          optimizer=optimizer,
          scheduler=scheduler,
          train_sampler=train_sampler,
          epoch=resume_from_epoch,
          batch=0,
      )
      state.register_reset_callbacks([on_state_reset])
  
      @hvd.elastic.run
      def full_train(state, train_loader):
          while state.epoch < args.epochs:
              train(state, train_loader)
              state.scheduler.step()
              save_checkpoint(state)
              end_epoch(state)
  
      full_train(state, train_loader)
  
imports: import argparse
  import os
  import numpy
  from tqdm import tqdm
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  import horovod.torch
  from horovod.torch.elastic.sampler import ElasticSampler
  from horovod.ray import ray_logger
  from horovod.ray.elastic import TestDiscovery
label_desc: Use image node_id 6 for calling a global function with example usage: run(large) without return types