def METHOD_NAME(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()