code: def train(
      logger,
      config,
      train_loader,
      model,
      optimizer,
      criterion,
      epoch,
      main_proc,
  ):
      meters = AverageMeterGroup()
      cur_lr = optimizer.param_groups[0]["lr"]
      if main_proc:
          logger.info("Epoch %d LR %.6f", epoch, cur_lr)
  
      model.train()
      for step, (x, y) in enumerate(train_loader):
          x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
          optimizer.zero_grad()
          logits, aux_logits = model(x)
          loss = criterion(logits, y)
          if config.aux_weight > 0.0:
              loss += config.aux_weight * criterion(aux_logits, y)
          loss.backward()
          nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
          optimizer.step()
  
          prec1, prec5 = utils.accuracy(logits, y, topk=(1, 5))
          metrics = {"prec1": prec1, "prec5": prec5, "loss": loss}
          metrics = utils.reduce_metrics(metrics, config.distributed)
          meters.update(metrics)
  
          if main_proc and (
              step % config.log_frequency == 0
              or step + 1 == len(train_loader)
          ):
              logger.info(
                  "Epoch [%d/%d] Step [%d/%d]  %s",
                  epoch + 1,
                  config.epochs,
                  step + 1,
                  len(train_loader),
                  meters,
              )
  
      if main_proc:
          logger.info(
              "Train: [%d/%d] Final Prec@1 %.4f Prec@5 %.4f",
              epoch + 1,
              config.epochs,
              meters.prec1.avg,
              meters.prec5.avg,
          )
  
file_path: nni/examples/nas/legacy/cdarts/retrain.py
imports: import json
  import logging
  import os
  import time
  from argparse import ArgumentParser
  import torch
  import torch.nn
  import apex
  import datasets
  import utils
  from apex.parallel import DistributedDataParallel
  from config import RetrainConfig
  from datasets.cifar import get_augment_datasets
  from model import Model
  from nni.nas.pytorch.fixed import apply_fixed_architecture
  from nni.nas.pytorch.utils import AverageMeterGroup
label_desc: Use image node_id 1 for calling a global function with example usage: train(logger, config, train_loader, model, optimizer, criterion, epoch, main_proc) without return types