def METHOD_NAME(state):
    start_batch = state.batch
    # Horovod: adjust number of steps based on number of GPUs.
    for batch_idx, (images, labels) in enumerate(dataset.skip(state.batch).take(100 // hvd.size())):
        state.batch = start_batch + batch_idx
        loss_value = training_step(images, labels)
        if state.batch % 10 == 0 and hvd.rank() == 0:
            print('Step #%d\tLoss: %.6f' % (state.batch, loss_value))
        # Horovod: commit state at the end of each batch
        state.commit()