code: def main_worker(gpu, ngpus_per_node, args):
      global best_acc1
      args.gpu = gpu
  
      if args.gpu is not None:
          print("Use GPU: {} for training".format(args.gpu))
  
      if args.pretrained:
          print("=> using pre-trained model '{}'".format(args.arch))
          model = models.__dict__[args.arch](pretrained=True)
      else:
          print("=> creating model '{}'".format(args.arch))
          model = models.__dict__[args.arch]()
  
      # In case of distributed process, initializes the distributed backend
      # which will take care of sychronizing nodes/GPUs
      if args.local_rank == -1:
          if args.gpu:
              device = torch.device("cuda:{}".format(args.gpu))
          else:
              device = torch.device("cuda")
      else:
          torch.cuda.set_device(args.local_rank)
          device = torch.device("cuda", args.local_rank)
          deepspeed.init_distributed()
  
      def print_rank_0(msg):
          if args.local_rank <= 0:
              print(msg)
  
      args.batch_size = int(args.batch_size / ngpus_per_node)
      if (
          not torch.cuda.is_available()
      ):  # and not torch.backends.mps.is_available():
          print("using CPU, this will be slow")
          device = torch.device("cpu")
          model = model.to(device)
  
      # define loss function (criterion), optimizer, and learning rate scheduler
      criterion = nn.CrossEntropyLoss().to(device)
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
      scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
      # optionally resume from a checkpoint
      if args.resume:
          if os.path.isfile(args.resume):
              print("=> loading checkpoint '{}'".format(args.resume))
              if args.gpu is None:
                  checkpoint = torch.load(args.resume)
              elif torch.cuda.is_available():
                  # Map model to be loaded to specified single gpu.
                  loc = "cuda:{}".format(args.gpu)
                  checkpoint = torch.load(args.resume, map_location=loc)
              args.start_epoch = checkpoint["epoch"]
              best_acc1 = checkpoint["best_acc1"]
              if args.gpu is not None:
                  # best_acc1 may be from a checkpoint from a different GPU
                  best_acc1 = best_acc1.to(args.gpu)
              model.load_state_dict(checkpoint["state_dict"])
              optimizer.load_state_dict(checkpoint["optimizer"])
              scheduler.load_state_dict(checkpoint["scheduler"])
              print(
                  "=> loaded checkpoint '{}' (epoch {})".format(
                      args.resume, checkpoint["epoch"]
                  )
              )
          else:
              print(
                  "=> no checkpoint found at '{}'".format(args.resume)
              )
  
      # Initialize DeepSpeed for the model
      model, optimizer, _, _ = deepspeed.initialize(
          model=model,
          optimizer=optimizer,
          args=args,
          lr_scheduler=None,  # scheduler,
          dist_init_required=True,
      )
  
      # Data loading code
      if args.dummy:
          print("=> Dummy data is used!")
          train_dataset = datasets.FakeData(
              1281167, (3, 224, 224), 1000, transforms.ToTensor()
          )
          val_dataset = datasets.FakeData(
              50000, (3, 224, 224), 1000, transforms.ToTensor()
          )
      else:
          traindir = os.path.join(args.data, "train")
          valdir = os.path.join(args.data, "val")
          normalize = transforms.Normalize(
              mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
          )
  
          train_dataset = datasets.ImageFolder(
              traindir,
              transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
  
          val_dataset = datasets.ImageFolder(
              valdir,
              transforms.Compose(
                  [
                      transforms.Resize(256),
                      transforms.CenterCrop(224),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
  
      if args.local_rank != -1:
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  train_dataset
              )
          )
          val_sampler = torch.utils.data.distributed.DistributedSampler(
              val_dataset, shuffle=False, drop_last=True
          )
      else:
          train_sampler = None
          val_sampler = None
  
      print("Batch_size:", args.batch_size)
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=args.batch_size,
          shuffle=(train_sampler is None),
          num_workers=args.workers,
          pin_memory=True,
          sampler=train_sampler,
      )
  
      val_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.workers,
          pin_memory=True,
          sampler=val_sampler,
      )
  
      if args.evaluate:
          validate(val_loader, model, criterion, args)
          return
  
      losses = torch.empty(args.epochs).cuda()
      acc1s = torch.empty(args.epochs).cuda()
      for epoch in range(args.start_epoch, args.epochs):
          if args.distributed:
              train_sampler.set_epoch(epoch)
          # train for one epoch
          this_loss = train(
              train_loader,
              model,
              criterion,
              optimizer,
              epoch,
              device,
              args,
          )
          losses[epoch] = this_loss
  
          # evaluate on validation set
          acc1 = validate(val_loader, model, criterion, args)
          acc1s[epoch] = acc1
  
          scheduler.step()
  
          # remember best acc@1 and save checkpoint
          is_best = acc1 > best_acc1
          best_acc1 = max(acc1, best_acc1)
  
          if not args.multiprocessing_distributed or (
              args.multiprocessing_distributed and args.gpu is None
          ):
              save_checkpoint(
                  {
                      "epoch": epoch + 1,
                      "arch": args.arch,
                      "state_dict": model.state_dict(),
                      "best_acc1": best_acc1,
                      "optimizer": optimizer.state_dict(),
                      "scheduler": scheduler.state_dict(),
                  },
                  is_best,
              )
  
      return (losses, acc1s)
  
imports: import argparse
  import os
  import pdb
  import random
  import shutil
  import time
  import warnings
  from enum import Enum
  import torch
  import deepspeed
  from openpyxl import Workbook
  import torch.backends.cudnn
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.nn.parallel
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  import torchvision.datasets
  import torchvision.models
  import torchvision.transforms
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.data import Subset
label_desc: Use image node_id 2 for calling a global function with example usage: main_worker(gpu, ngpus_per_node, args) and returns: losses,  acc1s