def for_and_backward(model, batch, maps, optimizer, loss_fn):
    """Does one forward and one backward step, including weight updates, using
    the given batch. Returns the loss."""
    prediction = model(batch)
    loss = loss_fn(prediction, maps)
    # Calculate gradients
    loss.backward()
    # Update model parameters from gradients
    optimizer.step()
    # Reset gradients
    optimizer.zero_grad()

    return loss