code: def training_step(
      batch: Any, model: torch.nn.Module, *args, **kwargs
  ):
      output = model(batch[0])
      loss = F.cross_entropy(output, batch[1])
      return loss
  
file_path: nni/examples/compression/evaluator/torch_evaluator.py
imports: from typing import Any, Callable
  import torch
  from examples.compression.models import prepare_dataloader
  import nni
  from examples.compression.models import build_resnet18
  import torch.nn.functional
  from nni.compression import TorchEvaluator
label_desc: Use image node_id 2 for calling a global function with example usage: training_step(batch, model) and returns: loss