import torch


def train_batch(model, batch):
    """Perform one iteration of model training given a single batch"""
    # send data to CUDA if necessary
    inputs, correct_outputs = batch
    if model.device:
        inputs = inputs.to(model.device)
        correct_outputs = correct_outputs.to(model.device)

    # train batch
    model_outputs = model.forward(inputs)
    model.optimizer.zero_grad()
    loss = model.criterion(model_outputs, correct_outputs)
    loss.backward(retain_graph=True)
    model.optimizer.step()

    # return data to CPU if necessary
    if model.device:
        inputs = inputs.to(torch.device("cpu"))
        correct_outputs = correct_outputs.to(torch.device("cpu"))
    return float(loss) / model_outputs.shape[0]