def METHOD_NAME(batch, model):
    try:
        batch_ = next(batch)
    except BaseException:
        batch_ = batch
    tokens, types, labels, attention_mask = process_batch(batch_)
    # Forward model.
    args = get_args()
    output_tensor = model(tokens, attention_mask, tokentype_ids=types)
    return output_tensor, partial(loss_func, output_predictions, labels)