code: def train(config, train_loader, model, optimizer, criterion, epoch):
      top1 = AverageMeter("top1")
      top5 = AverageMeter("top5")
      losses = AverageMeter("losses")
  
      cur_step = epoch * len(train_loader)
      cur_lr = optimizer.param_groups[0]["lr"]
      logger.info("Epoch %d LR %.6f", epoch, cur_lr)
      writer.add_scalar("lr", cur_lr, global_step=cur_step)
  
      model.train()
  
      for step, (x, y) in enumerate(train_loader):
          x, y = x.to(device, non_blocking=True), y.to(
              device, non_blocking=True
          )
          bs = x.size(0)
  
          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()
          # gradient clipping
          nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
          optimizer.step()
  
          accuracy = utils.accuracy(logits, y, topk=(1, 5))
          losses.update(loss.item(), bs)
          top1.update(accuracy["acc1"], bs)
          top5.update(accuracy["acc5"], bs)
          writer.add_scalar(
              "loss/train", loss.item(), global_step=cur_step
          )
          writer.add_scalar(
              "acc1/train", accuracy["acc1"], global_step=cur_step
          )
          writer.add_scalar(
              "acc5/train", accuracy["acc5"], global_step=cur_step
          )
  
          if (
              step % config.log_frequency == 0
              or step == len(train_loader) - 1
          ):
              logger.info(
                  "Train: [{:3d}/{}] Step {:03d}/{:03d} Loss {losses.avg:.3f} "
                  "Prec@(1,5) ({top1.avg:.1%}, {top5.avg:.1%})".format(
                      epoch + 1,
                      config.epochs,
                      step,
                      len(train_loader) - 1,
                      losses=losses,
                      top1=top1,
                      top5=top5,
                  )
              )
  
          cur_step += 1
  
      logger.info(
          "Train: [{:3d}/{}] Final Prec@1 {:.4%}".format(
              epoch + 1, config.epochs, top1.avg
          )
      )
  
file_path: nni/examples/nas/legacy/oneshot/darts/retrain.py
imports: import logging
  import time
  from argparse import ArgumentParser
  import torch
  import torch.nn
  from torch.utils.tensorboard import SummaryWriter
  import datasets
  import utils
  from model import CNN
  from nni.nas.pytorch.utils import AverageMeter
  from nni.retiarii import fixed_arch
label_desc: Use image node_id 1 for calling a global function with example usage: train(config, train_loader, model, optimizer, criterion, epoch) without return types