import torch


def synthetic_loss_no_mask(x, reconstructed_x, mean, log_var, args):
    """
    MSE loss for reconstruction, KLD loss as per VAE. 
    Also want to output dimension (element) wise RCL and KLD
    """

    loss = torch.nn.MSELoss(size_average=False)
    loss_per_element = torch.nn.MSELoss(size_average=False, reduce=False)

    RCL = loss(reconstructed_x, x)
    RCL_per_element = loss_per_element(reconstructed_x, x)
    KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())
    KLD_per_element = -0.5 * (1 + log_var - mean.pow(2) - log_var.exp())

    return RCL + args.beta_vae*KLD, RCL, KLD, RCL_per_element, KLD_per_element