def METHOD_NAME(state):
    optimizer.zero_grad()
    output = model(data)
    loss = F.cross_entropy(output, target)
    loss.backward()
    optimizer.step()
    state.batch += 1
    if state.batch == args.num_batches_per_commit:
        state.batch = 0
        state.commit()