code: def main(args):
      args.total_iterations = sum(args.dataset_steps)
  
      # initialize DDP setting
      utils.setup_ddp(args)
      print(args)
  
      args.test_only = args.train_datasets is None
  
      # set the appropriate devices
      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)
  
      # select model architecture
      model = torchvision.prototype.models.depth.stereo.__dict__[
          args.model
      ](weights=args.weights)
  
      # convert to DDP if need be
      if args.distributed:
          model = model.to(args.gpu)
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.gpu]
          )
          model_without_ddp = model.module
      else:
          model.to(device)
          model_without_ddp = model
  
      os.makedirs(args.checkpoint_dir, exist_ok=True)
  
      val_loaders = {
          name: make_eval_loader(name, args)
          for name in args.test_datasets
      }
  
      # EVAL ONLY configurations
      if args.test_only:
          evaluate(model, val_loaders, args)
          return
  
      # Sanity check for the parameter count
      print(
          f"Parameter Count: {sum(p.numel() for p in model.parameters() if p.requires_grad)}"
      )
  
      # Compose the training dataset
      train_dataset = get_train_dataset(args.dataset_root, args)
  
      # initialize the optimizer
      if args.optimizer == "adam":
          optimizer = torch.optim.Adam(
              model.parameters(),
              lr=args.lr,
              weight_decay=args.weight_decay,
          )
      elif args.optimizer == "sgd":
          optimizer = torch.optim.SGD(
              model.parameters(),
              lr=args.lr,
              weight_decay=args.weight_decay,
              momentum=0.9,
          )
      else:
          raise ValueError(
              f"Unknown optimizer {args.optimizer}. Please choose between adam and sgd"
          )
  
      # initialize the learning rate schedule
      scheduler = make_lr_schedule(args, optimizer)
  
      # load them from checkpoint if needed
      args.start_step = 0
      if args.resume_path is not None:
          checkpoint = torch.load(
              args.resume_path, map_location="cpu", weights_only=True
          )
          if "model" in checkpoint:
              # this means the user requested to resume from a training checkpoint
              model_without_ddp.load_state_dict(checkpoint["model"])
              # this means the user wants to continue training from where it was left off
              if args.resume_schedule:
                  optimizer.load_state_dict(checkpoint["optimizer"])
                  scheduler.load_state_dict(checkpoint["scheduler"])
                  args.start_step = checkpoint["step"] + 1
                  # modify starting point of the dat
                  sample_start_step = (
                      args.start_step
                      * args.batch_size
                      * args.world_size
                  )
                  train_dataset = train_dataset[sample_start_step:]
  
          else:
              # this means the user wants to finetune on top of a model state dict
              # and that no other changes are required
              model_without_ddp.load_state_dict(checkpoint)
  
      torch.backends.cudnn.benchmark = True
  
      # enable training mode
      model.train()
      if args.freeze_batch_norm:
          freeze_batch_norm(model_without_ddp)
  
      # put dataloader on top of the dataset
      # make sure to disable shuffling since the dataset is already shuffled
      # in order to guarantee quasi randomness whilst retaining a deterministic
      # dataset consumption order
      if args.distributed:
          # the train dataset is preshuffled in order to respect the iteration order
          sampler = torch.utils.data.distributed.DistributedSampler(
              train_dataset, shuffle=False, drop_last=True
          )
      else:
          # the train dataset is already shuffled, so we can use a simple SequentialSampler
          sampler = torch.utils.data.SequentialSampler(train_dataset)
  
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          sampler=sampler,
          batch_size=args.batch_size,
          pin_memory=True,
          num_workers=args.workers,
      )
  
      # initialize the logger
      if args.tensorboard_summaries:
          from torch.utils.tensorboard import SummaryWriter
  
          tensorboard_path = Path(args.checkpoint_dir) / "tensorboard"
          os.makedirs(tensorboard_path, exist_ok=True)
  
          tensorboard_run = tensorboard_path / f"{args.name}"
          writer = SummaryWriter(tensorboard_run)
      else:
          writer = None
  
      logger = utils.MetricLogger(delimiter="  ")
  
      scaler = (
          torch.cuda.amp.GradScaler() if args.mixed_precision else None
      )
      # run the training loop
      # this will perform optimization, respectively logging and saving checkpoints
      # when need be
      run(
          model=model,
          optimizer=optimizer,
          scheduler=scheduler,
          train_loader=train_loader,
          val_loaders=val_loaders,
          logger=logger,
          writer=writer,
          scaler=scaler,
          args=args,
      )
  
imports: import argparse
  import os
  import warnings
  from pathlib import Path
  from typing import List, Union
  import numpy
  import torch
  import torch.distributed
  import torchvision.models.optical_flow
  import torchvision.prototype.models.depth.stereo
  import utils
  import visualization
  from parsing import make_dataset, make_eval_transform, make_train_transform, VALID_DATASETS
  from torch import nn
  from torchvision.transforms.functional import get_dimensions, InterpolationMode, resize
  from utils.metrics import AVAILABLE_METRICS
  from utils.norm import freeze_batch_norm
label_desc: Use image node_id 10 for calling a global function with example usage: main(args) without return types