def train_supervised(
    model,
    input_tensor,
    y_true,
    loss_fn,
    optimizer,
    multiclass =False,
    n_out = 1,
    ):
    """
    Helper function to make forward and backward pass with minibatch.
    
    Params
    ------
    n_out (int, default = 1)
        Dimensionality of output dimension. Leave as 1 for multiclass, 
        i.e. the output is a probability distribution over classes (e.g. MNIST).
    """
    
    # Zero out grads 
    model.zero_grad()
    y_pred = model(input_tensor)
    
    #Note that if it's a multiclass classification (i.e. the output is a 
    # probability distribution over classes) the loss_fn 
    # nn.NLLLoss(y_pred, y_true) uses as input y_pred.size = (n_batch, n_classes)
    # and y_true.size = (n_batch), that's why it doesn't get reshaped. 

    if multiclass: 
        loss = loss_fn(y_pred, y_true)

    else: # Backprop error
        loss = loss_fn(y_pred, y_true.view(-1, n_out).float())
    
    loss.backward()
    # Update weights 
    optimizer.step()
    
    return loss