code: def load_checkpoint(args):
      utils.setup_ddp(args)
  
      if not args.weights:
          checkpoint = torch.load(
              args.checkpoint,
              map_location=torch.device("cpu"),
              weights_only=True,
          )
          if "model" in checkpoint:
              experiment_args = checkpoint["args"]
              model = (
                  torchvision.prototype.models.depth.stereo.__dict__[
                      experiment_args.model
                  ](weights=None)
              )
              model.load_state_dict(checkpoint["model"])
          else:
              model = (
                  torchvision.prototype.models.depth.stereo.__dict__[
                      args.model
                  ](weights=None)
              )
              model.load_state_dict(checkpoint)
  
          # 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)
      else:
          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.device)
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.gpu]
          )
      else:
          model.to(device)
  
      return model
  
file_path: vision/references/depth/stereo/cascade_evaluation.py
imports: import os
  import warnings
  import torch
  import torchvision
  import torchvision.prototype.models.depth.stereo
  import utils
  from torch.nn import functional
  from train import make_eval_loader
  from utils.metrics import AVAILABLE_METRICS
  from visualization import make_prediction_image_side_to_side
label_desc: Use image node_id 5 for calling a global function with example usage: load_checkpoint(args) and returns: model