code: def training_step(batch, model: torch.nn.Module):
      output = model(batch[0])
      loss = F.cross_entropy(output, batch[1])
      return loss
  
file_path: nni/examples/compression/models.py
imports: from pathlib import Path
  import torch
  import torch.nn.functional
  from torch.optim import Adam
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from torchvision import datasets, transforms
  from torchvision.models.mobilenetv3 import mobilenet_v3_small
  from torchvision.models.resnet import resnet18
  import nni
label_desc: Use image node_id 7 for calling a global function with example usage: training_step(batch, model) and returns: loss