def METHOD_NAME(state):
    state.rendezvous += 1
    while state.epoch < args.epochs:
        print('epoch {} batch {}'.format(state.epoch, state.batch))
        while state.batch < args.batches_per_epoch:
            check_exit(state.epoch, state.batch)
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
            state.batch += 1
            if state.batch % args.batches_per_commit == 0:
                state.commits += 1
                state.commit()
        if hvd.rank() == 0:
            log_state(state)
            current_hosts = epoch_to_hosts.get(state.epoch, default_hosts)
            next_hosts = epoch_to_hosts.get(state.epoch + 1, default_hosts)
            if args.discovery_wait > 0 and current_hosts != next_hosts:
                print('host changes: {} -> {}'.format(current_hosts, next_hosts))
                start = int(time.time())
                # Reset the last updated timestamp to ensure we still raise HostsUpdatedInterrupt even
                # if the timestamp from the driver isn't greater than the last (due to time loops).
                state._last_updated_timestamp = 0
                while state._host_messages.empty():
                    if int(time.time()) - start > args.discovery_wait:
                        raise TimeoutError('Timed out waiting for notifications from driver.')
                    time.sleep(0.1)
        if args.epoch_wait > 0:
            time.sleep(args.epoch_wait)
        state.epoch += 1
        state.batch = 0
        state.commits += 1
        state.commit()