def update(model, optimizer, train_x, train_y, ret_loss=False):
    """Perform a single model update 

    Apply model on the given data train_x, train_y.
    Compute the prediction loss and make parameter updates using the provided
    optimizer.

    Parameters
    ----------
    model : nn.Module
        Pytorch neural network module. Must have a criterion attribute!!
    optimizer : torch.optim
        Optimizer that updates the model
    train_x : torch.Tensor
        Input tensor x
    train_y : torch.Tensor 
        Ground-truth output tensor y
    ret_loss : bool, optional
        Whether to return the observed loss (default=False)
        
    Returns
    ------
    ret_loss (optional) 
        Only if ret_loss=True
    """
    
    model.zero_grad()
    out = model(train_x)
    loss = model.criterion(out, train_y)
    loss.backward()
    optimizer.step()
    if ret_loss:
        return loss.item()