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.distributed:
          if args.dist_url == "env://" and args.rank == -1:
              args.rank = int(os.environ["RANK"])
          if args.multiprocessing_distributed:
              # For multiprocessing distributed training, rank needs to be the
              # global rank among all the processes
              args.rank = args.rank * ngpus_per_node + gpu
          dist.init_process_group(
              backend=args.dist_backend,
              init_method=args.dist_url,
              world_size=args.world_size,
              rank=args.rank,
          )
  
      # create model
      print("=> creating model '{}'".format(args.arch))
      if args.arch.split("_")[-1][:-1] == "lpf":  # antialiased model
          model = antialiased_cnns.__dict__[args.arch[:-5]](
              pretrained=args.pretrained,
              filter_size=int(args.arch[-1]),
              _force_nonfinetuned=args.force_nonfinetuned,
          )
      else:  # baseline model
          model = models.__dict__[args.arch](pretrained=args.pretrained)
  
      # instrumentation
      if args.wandb:
          import wandb
  
          wandb.init(project="antialiased-cnns")
          wandb.config.update(args)
          wandb.watch(model)
  
      if args.finetune:  # finetune from baseline "aliased" model
          print(
              "=> copying over pretrained weights from [%s]"
              % args.arch[:-5]
          )
          model_baseline = models.__dict__[args.arch[:-5]](
              pretrained=True
          )
          antialiased_cnns.copy_params_buffers(model_baseline, model)
  
      if args.weights is not None:
          print("=> using saved weights [%s]" % args.weights)
          weights = torch.load(args.weights)
          model.load_state_dict(weights["state_dict"])
  
      if args.distributed:
          # For multiprocessing distributed, DistributedDataParallel constructor
          # should always set the single device scope, otherwise,
          # DistributedDataParallel will use all available devices.
          if args.gpu is not None:
              torch.cuda.set_device(args.gpu)
              model.cuda(args.gpu)
              # When using a single GPU per process and per
              # DistributedDataParallel, we need to divide the batch size
              # ourselves based on the total number of GPUs we have
              args.batch_size = int(args.batch_size / ngpus_per_node)
              args.workers = int(args.workers / ngpus_per_node)
              model = torch.nn.parallel.DistributedDataParallel(
                  model, device_ids=[args.gpu]
              )
          else:
              model.cuda()
              # DistributedDataParallel will divide and allocate batch_size to all
              # available GPUs if device_ids are not set
              model = torch.nn.parallel.DistributedDataParallel(model)
      elif args.gpu is not None:
          torch.cuda.set_device(args.gpu)
          model = model.cuda(args.gpu)
      else:
          # DataParallel will divide and allocate batch_size to all available GPUs
          if args.arch.startswith("alexnet") or args.arch.startswith(
              "vgg"
          ):
              model.features = torch.nn.DataParallel(model.features)
              model.cuda()
          else:
              model = torch.nn.DataParallel(model).cuda()
  
      # define loss function (criterion) and optimizer
      criterion = nn.CrossEntropyLoss().cuda(args.gpu)
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      # optionally resume from a checkpoint
      if args.resume:
          if os.path.isfile(args.resume):
              print("=> loading checkpoint '{}'".format(args.resume))
              checkpoint = torch.load(args.resume)
              model.load_state_dict(
                  checkpoint["state_dict"], strict=False
              )
              if (
                  "optimizer" in checkpoint.keys()
              ):  # if no optimizer, then only load weights
                  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)
                  optimizer.load_state_dict(checkpoint["optimizer"])
              else:
                  print("  No optimizer saved")
              print(
                  "=> loaded checkpoint '{}' (epoch {})".format(
                      args.resume, checkpoint["epoch"]
                  )
              )
          else:
              print(
                  "=> no checkpoint found at '{}'".format(args.resume)
              )
  
      cudnn.benchmark = True
  
      # Data loading code
      traindir = os.path.join(args.data, "train")
      valdir = os.path.join(args.data, "val")
      mean = [0.485, 0.456, 0.406]
      std = [0.229, 0.224, 0.225]
      normalize = transforms.Normalize(mean=mean, std=std)
  
      if args.no_data_aug:
          train_dataset = datasets.ImageFolder(
              traindir,
              transforms.Compose(
                  [
                      transforms.Resize(256),
                      transforms.CenterCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
      else:
          train_dataset = datasets.ImageFolder(
              traindir,
              transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
  
      if args.distributed:
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  train_dataset
              )
          )
      else:
          train_sampler = None
  
      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,
      )
  
      crop_size = (
          256
          if (
              args.evaluate_shift
              or args.evaluate_diagonal
              or args.evaluate_save
          )
          else 224
      )
      args.batch_size = (
          1
          if (args.evaluate_diagonal or args.evaluate_save)
          else args.batch_size
      )
  
      val_loader = torch.utils.data.DataLoader(
          datasets.ImageFolder(
              valdir,
              transforms.Compose(
                  [
                      transforms.Resize(256),
                      transforms.CenterCrop(crop_size),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          ),
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.workers,
          pin_memory=True,
      )
  
      if (
          args.val_debug
      ):  # debug mode - train on val set for faster epochs
          train_loader = val_loader
  
      if args.embed:
          from IPython import embed
  
          embed()
  
      if args.save_weights is not None:  # "deparallelize" saved weights
          print(
              "=> saving 'deparallelized' weights [%s]"
              % args.save_weights
          )
          # TO-DO: automatically save this during training
          if args.gpu is not None:
              torch.save(
                  {"state_dict": model.state_dict()},
                  args.save_weights,
                  _use_new_zipfile_serialization=False,
              )
          else:
              if args.arch[:7] == "alexnet" or args.arch[:3] == "vgg":
                  model.features = model.features.module
                  torch.save(
                      {"state_dict": model.state_dict()},
                      args.save_weights,
                      _use_new_zipfile_serialization=False,
                  )
              else:
                  torch.save(
                      {"state_dict": model.module.state_dict()},
                      args.save_weights,
                      _use_new_zipfile_serialization=False,
                  )
          return
  
      if args.evaluate:
          validate(val_loader, model, criterion, args)
          return
  
      if args.evaluate_shift:
          validate_shift(val_loader, model, args)
          return
  
      if args.evaluate_diagonal:
          validate_diagonal(val_loader, model, args)
          return
  
      if args.evaluate_save:
          validate_save(val_loader, mean, std, args)
          return
  
      if args.cos_lr:
          scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
              optimizer, args.epochs
          )
          for epoch in range(args.start_epoch):
              scheduler.step()
  
      for epoch in range(args.start_epoch, args.epochs):
          if args.distributed:
              train_sampler.set_epoch(epoch)
  
          if not args.cos_lr:
              adjust_learning_rate(optimizer, epoch, args)
          else:
              scheduler.step()
              print("[%03d] %.5f" % (epoch, scheduler.get_lr()[0]))
  
          if args.wandb:
              wandb.log(
                  {"learning_rate": optimizer.param_groups[0]["lr"]},
                  commit=False,
              )
  
          # train for one epoch
          train(train_loader, model, criterion, optimizer, epoch, args)
  
          # evaluate on validation set
          acc1 = validate(val_loader, model, criterion, args)
  
          # 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.rank % ngpus_per_node == 0
          ):
              save_checkpoint(
                  {
                      "epoch": epoch + 1,
                      "arch": args.arch,
                      "state_dict": model.state_dict(),
                      "best_acc1": best_acc1,
                      "optimizer": optimizer.state_dict(),
                  },
                  is_best,
                  epoch,
                  out_dir=args.out_dir,
              )
  
imports: import argparse
  import os
  import random
  import shutil
  import time
  import warnings
  import sys
  import numpy
  import os
  import torch
  import torch.nn
  import torch.nn.parallel
  import torch.backends.cudnn
  import torch.distributed
  import torch.optim
  import torch.multiprocessing
  import torch.utils.data
  import torch.utils.data.distributed
  import torchvision.transforms
  import torchvision.datasets
  import antialiased_cnns
  import torchvision.models
label_desc: Use image node_id 2 for calling a global function with example usage: main_worker(gpu, ngpus_per_node, args) without return types