code: def main(args):
      if args.output_dir:
          utils.mkdir(args.output_dir)
  
      utils.init_distributed_mode(args)
      print(args)
  
      device = torch.device(args.device)
  
      if args.use_deterministic_algorithms:
          torch.backends.cudnn.benchmark = False
          torch.use_deterministic_algorithms(True)
      else:
          torch.backends.cudnn.benchmark = True
  
      train_dir = os.path.join(args.data_path, "train")
      val_dir = os.path.join(args.data_path, "val")
      dataset, dataset_test, train_sampler, test_sampler = load_data(
          train_dir, val_dir, args
      )
  
      num_classes = len(dataset.classes)
      mixup_cutmix = get_mixup_cutmix(
          mixup_alpha=args.mixup_alpha,
          cutmix_alpha=args.cutmix_alpha,
          num_categories=num_classes,
          use_v2=args.use_v2,
      )
      if mixup_cutmix is not None:
  
          def collate_fn(batch):
              return mixup_cutmix(*default_collate(batch))
  
      else:
          collate_fn = default_collate
  
      data_loader = torch.utils.data.DataLoader(
          dataset,
          batch_size=args.batch_size,
          sampler=train_sampler,
          num_workers=args.workers,
          pin_memory=True,
          collate_fn=collate_fn,
      )
      data_loader_test = torch.utils.data.DataLoader(
          dataset_test,
          batch_size=args.batch_size,
          sampler=test_sampler,
          num_workers=args.workers,
          pin_memory=True,
      )
  
      print("Creating model")
      model = torchvision.models.get_model(
          args.model, weights=args.weights, num_classes=num_classes
      )
      model.to(device)
  
      if args.distributed and args.sync_bn:
          model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
  
      criterion = nn.CrossEntropyLoss(
          label_smoothing=args.label_smoothing
      )
  
      custom_keys_weight_decay = []
      if args.bias_weight_decay is not None:
          custom_keys_weight_decay.append(
              ("bias", args.bias_weight_decay)
          )
      if args.transformer_embedding_decay is not None:
          for key in [
              "class_token",
              "position_embedding",
              "relative_position_bias_table",
          ]:
              custom_keys_weight_decay.append(
                  (key, args.transformer_embedding_decay)
              )
      parameters = utils.set_weight_decay(
          model,
          args.weight_decay,
          norm_weight_decay=args.norm_weight_decay,
          custom_keys_weight_decay=custom_keys_weight_decay
          if len(custom_keys_weight_decay) > 0
          else None,
      )
  
      opt_name = args.opt.lower()
      if opt_name.startswith("sgd"):
          optimizer = torch.optim.SGD(
              parameters,
              lr=args.lr,
              momentum=args.momentum,
              weight_decay=args.weight_decay,
              nesterov="nesterov" in opt_name,
          )
      elif opt_name == "rmsprop":
          optimizer = torch.optim.RMSprop(
              parameters,
              lr=args.lr,
              momentum=args.momentum,
              weight_decay=args.weight_decay,
              eps=0.0316,
              alpha=0.9,
          )
      elif opt_name == "adamw":
          optimizer = torch.optim.AdamW(
              parameters, lr=args.lr, weight_decay=args.weight_decay
          )
      else:
          raise RuntimeError(
              f"Invalid optimizer {args.opt}. Only SGD, RMSprop and AdamW are supported."
          )
  
      scaler = torch.cuda.amp.GradScaler() if args.amp else None
  
      args.lr_scheduler = args.lr_scheduler.lower()
      if args.lr_scheduler == "steplr":
          main_lr_scheduler = torch.optim.lr_scheduler.StepLR(
              optimizer,
              step_size=args.lr_step_size,
              gamma=args.lr_gamma,
          )
      elif args.lr_scheduler == "cosineannealinglr":
          main_lr_scheduler = (
              torch.optim.lr_scheduler.CosineAnnealingLR(
                  optimizer,
                  T_max=args.epochs - args.lr_warmup_epochs,
                  eta_min=args.lr_min,
              )
          )
      elif args.lr_scheduler == "exponentiallr":
          main_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
              optimizer, gamma=args.lr_gamma
          )
      else:
          raise RuntimeError(
              f"Invalid lr scheduler '{args.lr_scheduler}'. Only StepLR, CosineAnnealingLR and ExponentialLR "
              "are supported."
          )
  
      if args.lr_warmup_epochs > 0:
          if args.lr_warmup_method == "linear":
              warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
                  optimizer,
                  start_factor=args.lr_warmup_decay,
                  total_iters=args.lr_warmup_epochs,
              )
          elif args.lr_warmup_method == "constant":
              warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
                  optimizer,
                  factor=args.lr_warmup_decay,
                  total_iters=args.lr_warmup_epochs,
              )
          else:
              raise RuntimeError(
                  f"Invalid warmup lr method '{args.lr_warmup_method}'. Only linear and constant are supported."
              )
          lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
              optimizer,
              schedulers=[warmup_lr_scheduler, main_lr_scheduler],
              milestones=[args.lr_warmup_epochs],
          )
      else:
          lr_scheduler = main_lr_scheduler
  
      model_without_ddp = model
      if args.distributed:
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.gpu]
          )
          model_without_ddp = model.module
  
      model_ema = None
      if args.model_ema:
          # Decay adjustment that aims to keep the decay independent of other hyper-parameters originally proposed at:
          # https://github.com/facebookresearch/pycls/blob/f8cd9627/pycls/core/net.py#L123
          #
          # total_ema_updates = (Dataset_size / n_GPUs) * epochs / (batch_size_per_gpu * EMA_steps)
          # We consider constant = Dataset_size for a given dataset/setup and omit it. Thus:
          # adjust = 1 / total_ema_updates ~= n_GPUs * batch_size_per_gpu * EMA_steps / epochs
          adjust = (
              args.world_size
              * args.batch_size
              * args.model_ema_steps
              / args.epochs
          )
          alpha = 1.0 - args.model_ema_decay
          alpha = min(1.0, alpha * adjust)
          model_ema = utils.ExponentialMovingAverage(
              model_without_ddp, device=device, decay=1.0 - alpha
          )
  
      if args.resume:
          checkpoint = torch.load(
              args.resume, map_location="cpu", weights_only=True
          )
          model_without_ddp.load_state_dict(checkpoint["model"])
          if not args.test_only:
              optimizer.load_state_dict(checkpoint["optimizer"])
              lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
          args.start_epoch = checkpoint["epoch"] + 1
          if model_ema:
              model_ema.load_state_dict(checkpoint["model_ema"])
          if scaler:
              scaler.load_state_dict(checkpoint["scaler"])
  
      if args.test_only:
          # We disable the cudnn benchmarking because it can noticeably affect the accuracy
          torch.backends.cudnn.benchmark = False
          torch.backends.cudnn.deterministic = True
          if model_ema:
              evaluate(
                  model_ema,
                  criterion,
                  data_loader_test,
                  device=device,
                  log_suffix="EMA",
              )
          else:
              evaluate(
                  model, criterion, data_loader_test, device=device
              )
          return
  
      print("Start training")
      start_time = time.time()
      for epoch in range(args.start_epoch, args.epochs):
          if args.distributed:
              train_sampler.set_epoch(epoch)
          train_one_epoch(
              model,
              criterion,
              optimizer,
              data_loader,
              device,
              epoch,
              args,
              model_ema,
              scaler,
          )
          lr_scheduler.step()
          evaluate(model, criterion, data_loader_test, device=device)
          if model_ema:
              evaluate(
                  model_ema,
                  criterion,
                  data_loader_test,
                  device=device,
                  log_suffix="EMA",
              )
          if args.output_dir:
              checkpoint = {
                  "model": model_without_ddp.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "lr_scheduler": lr_scheduler.state_dict(),
                  "epoch": epoch,
                  "args": args,
              }
              if model_ema:
                  checkpoint["model_ema"] = model_ema.state_dict()
              if scaler:
                  checkpoint["scaler"] = scaler.state_dict()
              utils.save_on_master(
                  checkpoint,
                  os.path.join(args.output_dir, f"model_{epoch}.pth"),
              )
              utils.save_on_master(
                  checkpoint,
                  os.path.join(args.output_dir, "checkpoint.pth"),
              )
  
      total_time = time.time() - start_time
      total_time_str = str(datetime.timedelta(seconds=int(total_time)))
      print(f"Training time {total_time_str}")
  
imports: import datetime
  import os
  import time
  import warnings
  import presets
  import torch
  import torch.utils.data
  import torchvision
  import torchvision.transforms
  import utils
  from sampler import RASampler
  from torch import nn
  from torch.utils.data.dataloader import default_collate
  from torchvision.transforms.functional import InterpolationMode
  from transforms import get_mixup_cutmix
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) and returns: mixup_cutmix