code: def main(args, config):
      utils.init_distributed_mode(args)
  
      device = torch.device(args.device)
  
      # fix the seed for reproducibility
      seed = args.seed + utils.get_rank()
      torch.manual_seed(seed)
      np.random.seed(seed)
      random.seed(seed)
      cudnn.benchmark = True
  
      #### Dataset ####
      print("Creating dataset")
      datasets = [create_dataset("pretrain", config, min_scale=0.2)]
      print("number of training samples: %d" % len(datasets[0]))
  
      num_tasks = utils.get_world_size()
      global_rank = utils.get_rank()
      samplers = create_sampler(
          datasets, [True], num_tasks, global_rank
      )
  
      data_loader = create_loader(
          datasets,
          samplers,
          batch_size=[config["batch_size"]],
          num_workers=[4],
          is_trains=[True],
          collate_fns=[None],
      )[0]
  
      #### Model ####
      print("Creating model")
      model = blip_pretrain(
          image_size=config["image_size"],
          vit=config["vit"],
          vit_grad_ckpt=config["vit_grad_ckpt"],
          vit_ckpt_layer=config["vit_ckpt_layer"],
          queue_size=config["queue_size"],
      )
  
      model = model.to(device)
  
      optimizer = torch.optim.AdamW(
          params=model.parameters(),
          lr=config["init_lr"],
          weight_decay=config["weight_decay"],
      )
  
      start_epoch = 0
      if args.checkpoint:
          checkpoint = torch.load(args.checkpoint, map_location="cpu")
          state_dict = checkpoint["model"]
          model.load_state_dict(state_dict)
  
          optimizer.load_state_dict(checkpoint["optimizer"])
          start_epoch = checkpoint["epoch"] + 1
          print("resume checkpoint from %s" % args.checkpoint)
  
      model_without_ddp = model
      if args.distributed:
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.gpu]
          )
          model_without_ddp = model.module
  
      print("Start training")
      start_time = time.time()
      for epoch in range(start_epoch, config["max_epoch"]):
          step_lr_schedule(
              optimizer,
              epoch,
              config["init_lr"],
              config["min_lr"],
              config["lr_decay_rate"],
          )
  
          train_stats = train(
              model, data_loader, optimizer, epoch, device, config
          )
          if utils.is_main_process():
              log_stats = {
                  **{f"train_{k}": v for k, v in train_stats.items()},
                  "epoch": epoch,
              }
              save_obj = {
                  "model": model_without_ddp.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "config": config,
                  "epoch": epoch,
              }
              torch.save(
                  save_obj,
                  os.path.join(
                      args.output_dir, "checkpoint_%02d.pth" % epoch
                  ),
              )
  
              with open(
                  os.path.join(args.output_dir, "log.txt"), "a"
              ) as f:
                  f.write(json.dumps(log_stats) + "\n")
  
          dist.barrier()
  
      total_time = time.time() - start_time
      total_time_str = str(datetime.timedelta(seconds=int(total_time)))
      print("Training time {}".format(total_time_str))
  
imports: import argparse
  import os
  import ruamel_yaml
  import numpy
  import random
  import time
  import datetime
  import json
  from pathlib import Path
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.backends.cudnn
  import torch.distributed
  from torch.utils.data import DataLoader
  from models.blip_pretrain import blip_pretrain
  import utils
  from utils import warmup_lr_schedule, step_lr_schedule
  from data import create_dataset, create_sampler, create_loader
label_desc: Use image node_id 2 for calling a global function with example usage: main(args, config) without return types