def train_vae(model, input_tensor, opt, batch_size, loss_fn = None): 

    """
    Wrapper function to make a forward pass through a VAE model.
    Assumes that the model class contains a variational loss function, 
    if this is not the case it can be provided using the loss_fn arg. 
    
    Params 
    ------
    model (torch.nn.Module)
        VAE model. 

    input_tensor (torch.tensor)
        Tensor of size (batch_size, n_feats). 

    opt (torch.optim.Optimizer)

    batch_size (int)

    loss_fn (default= None)
        Loss function in case VAE class doesn't contain a loss function. 


    Returns 
    -------

    loss (float)
        Variational loss for the input batch. 

    """

    input_ = input_tensor.view(batch_size, -1 ).float()
    
    # Zero out grads 
    opt.zero_grad()

    # Make forward computation 
    reconstructed, mu, log_var = model(input_)

    # Backprop errors 
    loss = model.loss(reconstructed, input_, mu, log_var)
    loss.backward()

    # Update weights
    opt.step()

    return loss