code: def main(args):
      utils.setup_ddp(args)
      args.test_only = args.train_dataset is None
  
      if args.distributed and args.device == "cpu":
          raise ValueError(
              "The device must be cuda if we want to run in distributed mode using torchrun"
          )
      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
  
      model = torchvision.models.get_model(
          args.model, weights=args.weights
      )
  
      if args.distributed:
          model = model.to(args.local_rank)
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.local_rank]
          )
          model_without_ddp = model.module
      else:
          model.to(device)
          model_without_ddp = model
  
      if args.resume is not None:
          checkpoint = torch.load(
              args.resume, map_location="cpu", weights_only=True
          )
          model_without_ddp.load_state_dict(checkpoint["model"])
  
      if args.test_only:
          # Set deterministic CUDNN algorithms, since they can affect epe a fair bit.
          torch.backends.cudnn.benchmark = False
          torch.backends.cudnn.deterministic = True
          evaluate(model, args)
          return
  
      print(
          f"Parameter Count: {sum(p.numel() for p in model.parameters() if p.requires_grad)}"
      )
  
      train_dataset = get_train_dataset(
          args.train_dataset, args.dataset_root
      )
  
      optimizer = torch.optim.AdamW(
          model.parameters(),
          lr=args.lr,
          weight_decay=args.weight_decay,
          eps=args.adamw_eps,
      )
  
      scheduler = torch.optim.lr_scheduler.OneCycleLR(
          optimizer=optimizer,
          max_lr=args.lr,
          epochs=args.epochs,
          steps_per_epoch=ceil(
              len(train_dataset) / (args.world_size * args.batch_size)
          ),
          pct_start=0.05,
          cycle_momentum=False,
          anneal_strategy="linear",
      )
  
      if args.resume is not None:
          optimizer.load_state_dict(checkpoint["optimizer"])
          scheduler.load_state_dict(checkpoint["scheduler"])
          args.start_epoch = checkpoint["epoch"] + 1
      else:
          args.start_epoch = 0
  
      torch.backends.cudnn.benchmark = True
  
      model.train()
      if args.freeze_batch_norm:
          utils.freeze_batch_norm(model.module)
  
      if args.distributed:
          sampler = torch.utils.data.distributed.DistributedSampler(
              train_dataset, shuffle=True, drop_last=True
          )
      else:
          sampler = torch.utils.data.RandomSampler(train_dataset)
  
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          sampler=sampler,
          batch_size=args.batch_size,
          pin_memory=True,
          num_workers=args.workers,
      )
  
      logger = utils.MetricLogger()
  
      done = False
      for epoch in range(args.start_epoch, args.epochs):
          print(f"EPOCH {epoch}")
          if args.distributed:
              # needed on distributed mode, otherwise the data loading order would be the same for all epochs
              sampler.set_epoch(epoch)
  
          train_one_epoch(
              model=model,
              optimizer=optimizer,
              scheduler=scheduler,
              train_loader=train_loader,
              logger=logger,
              args=args,
          )
  
          # Note: we don't sync the SmoothedValues across processes, so the printed metrics are just those of rank 0
          print(f"Epoch {epoch} done. ", logger)
  
          if not args.distributed or args.rank == 0:
              checkpoint = {
                  "model": model_without_ddp.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "scheduler": scheduler.state_dict(),
                  "epoch": epoch,
                  "args": args,
              }
              torch.save(
                  checkpoint,
                  Path(args.output_dir) / f"{args.name}_{epoch}.pth",
              )
              torch.save(
                  checkpoint, Path(args.output_dir) / f"{args.name}.pth"
              )
  
          if epoch % args.val_freq == 0 or done:
              evaluate(model, args)
              model.train()
              if args.freeze_batch_norm:
                  utils.freeze_batch_norm(model.module)
  
file_path: vision/references/optical_flow/train.py
imports: import argparse
  import warnings
  from math import ceil
  from pathlib import Path
  import torch
  import torchvision.models.optical_flow
  import utils
  from presets import OpticalFlowPresetEval, OpticalFlowPresetTrain
  from torchvision.datasets import FlyingChairs, FlyingThings3D, HD1K, KittiFlow, Sintel
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types