def train_vae(model, train_loader, optimizer):
    """
    Function for training a model on a dataset. Train the model for one epoch.
    Inputs:
        model - VAE model to train
        train_loader - Data Loader for the dataset you want to train on
        optimizer - The optimizer used to update the parameters
    Outputs:
        average_bpd - Average BPD
        average_rec_loss - Average reconstruction loss
        average_reg_loss - Average regularization loss
    """

    # training loop
    bpd_vals, rec_loss, reg_loss = 0, 0, 0
    total_steps = 0
    for step, (train_inputs, train_targets) in enumerate(train_loader):

        # move to device
        train_inputs = train_inputs.to(model.device)

        # make a forward pass through the model (encoder and decoder)
        L_rec, L_reg, bpd = model.forward(train_inputs)

        # track bpd, rec, and reg
        bpd_vals += bpd.item()
        rec_loss += L_rec.item()
        reg_loss += L_reg.item()

        # backpropagation
        optimizer.zero_grad()
        bpd.backward()
        optimizer.step()

        # keep track of steps
        total_steps += 1

    average_bpd = bpd_vals / total_steps
    average_rec_loss = rec_loss / total_steps
    average_reg_loss = reg_loss / total_steps
    
    return average_bpd, average_rec_loss, average_reg_loss