def train_step(engine, batch):
    (x, y) = (batch[0], batch[1])
    if (x.device != device):
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
    model.train()
    y_pred = model(x)
    loss = criterion(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if ((config['log_every_iters'] > 0) and (((engine.state.iteration - 1) % config['log_every_iters']) == 0)):
        batch_loss = loss.item()
        engine.state.saved_batch_loss = batch_loss
    else:
        batch_loss = engine.state.saved_batch_loss
    return {'batch loss': batch_loss}
def train_step(engine, batch):
    (x, y) = (batch[0], batch[1])
    if (x.device != device):
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
    model.train()
    y_pred = model(x)
    loss = criterion(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if ((config['log_every_iters'] > 0) and (((engine.state.iteration - 1) % config['log_every_iters']) == 0)):
        batch_loss = loss.item()
        engine.state.saved_batch_loss = batch_loss
    else:
        batch_loss = engine.state.saved_batch_loss
    return {'batch loss': batch_loss}