code: def training_func(
      model: torch.nn.Module,
      optimizer: torch.optim.Optimizer,
      training_step: Callable[[Any, torch.nn.Module], torch.Tensor],
      lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
      max_steps: int,
      max_epochs: int,
  ):
      # create a train dataloader (and test dataloader if needs)
      train_dataloader, test_dataloader = prepare_dataloader()
  
      # deal with training duration, NNI prefers to prioritize the largest number of steps
      # at least `max_steps` or `max_epochs` will be given
      assert max_steps is not None or max_epochs is not None
      total_steps = (
          max_steps if max_steps else max_epochs * len(train_dataloader)
      )
      total_epochs = total_steps // len(train_dataloader) + (
          0 if total_steps % len(train_dataloader) == 0 else 1
      )
  
      # here is a common training loop
      current_step = 0
      for _ in range(total_epochs):
          for batch in train_dataloader:
              loss = training_step(batch, model)
              loss.backward()
              optimizer.step()
  
              # if reach the total steps, exit from the training loop
              current_step = current_step + 1
              if current_step >= total_steps:
                  return
  
          # if you are using a epoch-wise scheduler, call it here
          lr_scheduler.step()
  
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 1 for calling a global function with example usage: training_func(model, optimizer, training_step, lr_scheduler, max_steps, max_epochs) without return types