code: def train_epoch(
      model,
      auxiliarynet,
      criterion,
      train_loader,
      device,
      epoch,
      optimizer,
      logger,
  ):
      """Train one epoch."""
      model.train()
      auxiliarynet.train()
  
      batch_time = AverageMeter("batch_time")
      data_time = AverageMeter("data_time")
      losses = AverageMeter("losses")
  
      end = time.time()
      for i, (img, landmark_gt, angle_gt) in enumerate(train_loader):
          data_time.update(time.time() - end)
          img = img.to(device, non_blocking=True)
          landmark_gt = landmark_gt.to(device, non_blocking=True)
          angle_gt = angle_gt.to(device, non_blocking=True)
  
          lands, feats = model(img)
          landmarks = lands.squeeze()
          angle = auxiliarynet(feats)
  
          # task loss
          weighted_loss, _ = criterion(
              landmark_gt, angle_gt, angle, landmarks
          )
          loss = weighted_loss
  
          # compute gradient and do SGD step
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          # measure elapsed time
          batch_time.update(time.time() - end)
          end = time.time()
          # measure accuracy and record loss
          losses.update(
              np.squeeze(loss.cpu().detach().numpy()), img.size(0)
          )
  
          if i % 10 == 0:
              batch_log = (
                  "Train [{0}][{1}]\t"
                  "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
                  "Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
                  "Loss {losses.val:.4f} ({losses.avg:.4f})".format(
                      epoch + 1,
                      i,
                      batch_time=batch_time,
                      data_time=data_time,
                      losses=losses,
                  )
              )
              logger.info(batch_log)
  
file_path: nni/examples/nas/legacy/oneshot/pfld/retrain.py
imports: from __future__ import absolute_import, division, print_function
  import argparse
  import logging
  import os
  import time
  import torch
  import torchvision
  import numpy
  from datasets import PFLDDatasets
  from lib.builder import search_space
  from lib.ops import PRIMITIVES
  from lib.utils import PFLDLoss, accuracy
  from nni.algorithms.nas.pytorch.fbnet import LookUpTable, NASConfig, supernet_sample
  from nni.nas.pytorch.utils import AverageMeter
  from torch.utils.data import DataLoader
label_desc: Use image node_id 2 for calling a global function with example usage: train_epoch(model, auxiliarynet, criterion, train_loader, device, epoch, optimizer, logger) without return types