code: def train(
      model: torch.nn.Module,
      optimizer: torch.optim.Optimizer,
      training_step,
      lr_scheduler: _LRScheduler,
      max_steps: int,
      max_epochs: int,
  ):
      assert max_epochs is not None or max_steps is not None
      train_loader, test_loader = prepare_dataloader()
      max_steps = (
          max_steps if max_steps else max_epochs * len(train_loader)
      )
      max_epochs = max_steps // len(train_loader) + (
          0 if max_steps % len(train_loader) == 0 else 1
      )
      count_steps = 0
  
      model.train()
      for epoch in range(max_epochs):
          for data, target in train_loader:
              data, target = data.to(device), target.to(device)
              optimizer.zero_grad()
              loss = training_step((data, target), model)
              loss.backward()
              optimizer.step()
              count_steps += 1
              if count_steps >= max_steps:
                  acc = evaluate(model, test_loader)
                  print(
                      f"[Training Epoch {epoch} / Step {count_steps}] Final Acc: {acc}%"
                  )
                  return
          acc = evaluate(model, test_loader)
          print(
              f"[Training Epoch {epoch} / Step {count_steps}] Final Acc: {acc}%"
          )
  
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 5 for calling a global function with example usage: train(model, optimizer, training_step, lr_scheduler, max_steps, max_epochs) without return types