def METHOD_NAME():
    input_ids, label = next(train_dataloader)
    input_ids, label = input_ids.cuda(), label.cuda()
    with torch.no_grad():
        torch_output = torch_model(input_ids)
        torch_loss = criterion(torch_output.float(), label)
        zero_output = model(input_ids)
        zero_loss = criterion(zero_output.float(), label)
    assert_close(torch_loss, zero_loss)