code: def main(opts, is_iteration_based, **kwargs):
      device = getattr(opts, "dev.device", torch.device("cpu"))
  
      # set-up data loaders
      train_loader, val_loader, train_sampler = create_train_val_loader(
          opts
      )
  
      # compute max iterations based on max epochs
      # Useful in doing polynomial decay
      setattr(opts, "scheduler.is_iteration_based", is_iteration_based)
      if is_iteration_based:
          setattr(opts, "scheduler.max_iterations", 5)
          setattr(opts, "scheduler.max_epochs", 100000)
      else:
          setattr(opts, "scheduler.max_epochs", 2)
          setattr(opts, "scheduler.max_iterations", 10000)
  
      # set-up the model
      model = get_model(opts)
  
      # memory format
      memory_format = (
          torch.channels_last
          if getattr(opts, "common.channels_last", False)
          else torch.contiguous_format
      )
  
      model = model.to(device=device, memory_format=memory_format)
  
      # setup criteria
      criteria = build_loss_fn(opts)
      criteria = criteria.to(device=device)
  
      # create the optimizer
      optimizer = build_optimizer(model, opts=opts)
  
      # create the gradient scalar
      gradient_scaler = GradScaler(enabled=False)
  
      # LR scheduler
      scheduler = build_scheduler(opts=opts)
  
      model_ema = None
      use_ema = getattr(opts, "ema.enable", False)
  
      if use_ema:
          ema_momentum = getattr(opts, "ema.momentum", 0.0001)
          model_ema = EMA(
              model=model, ema_momentum=ema_momentum, device=device
          )
  
      best_metric = (
          0.0
          if getattr(opts, "stats.checkpoint_metric_max", False)
          else math.inf
      )
  
      start_epoch = 0
      start_iteration = 0
      resume_loc = getattr(opts, "common.resume", None)
      finetune_loc = getattr(opts, "common.finetune", None)
      auto_resume = getattr(opts, "common.auto_resume", False)
      if resume_loc is not None or auto_resume:
          (
              model,
              optimizer,
              gradient_scaler,
              start_epoch,
              start_iteration,
              best_metric,
              model_ema,
          ) = load_checkpoint(
              opts=opts,
              model=model,
              optimizer=optimizer,
              model_ema=model_ema,
              gradient_scaler=gradient_scaler,
          )
      elif finetune_loc is not None:
          model, model_ema = load_model_state(
              opts=opts, model=model, model_ema=model_ema
          )
  
      training_engine = Trainer(
          opts=opts,
          model=model,
          validation_loader=val_loader,
          training_loader=train_loader,
          optimizer=optimizer,
          criterion=criteria,
          scheduler=scheduler,
          start_epoch=start_epoch,
          start_iteration=start_iteration,
          best_metric=best_metric,
          model_ema=model_ema,
          gradient_scaler=gradient_scaler,
      )
  
      training_engine.run(train_sampler=train_sampler)
  
imports: import os.path
  import sys
  from pathlib import Path
  import pytest
  import math
  import random
  import shutil
  import torch
  from torch.cuda.amp import GradScaler
  from cvnets import EMA, get_model
  from engine import Trainer
  from loss_fn import build_loss_fn
  from optim import build_optimizer
  from optim.scheduler import build_scheduler
  from tests.configs import get_config
  from tests.dummy_loader import create_train_val_loader
  from tests.test_utils import unset_pretrained_models_from_opts
  from utils.checkpoint_utils import load_checkpoint, load_model_state
  from utils.common_utils import create_directories, device_setup
  from utils.ddp_utils import is_master
label_desc: Use image node_id 1 for calling a global function with example usage: main(opts, is_iteration_based) without return types