import torch


def vae_loss(x, recon, mu, logvar, beta):
    """Standard VAE Loss.

    Args:
        x: input image.
        recon: model output as reconstruction of input.
        mu: mean of latent.
        logvar: logvar of latent distribution.
        beta : Hyper-param in loss function.

    Returns:
        float: Loss over minibatches.
    """
    mse_loss = torch.nn.MSELoss().cuda()
    recon = recon.cuda()
    x = x.cuda()
    MSE = mse_loss(recon, x)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    KLD = KLD * beta
    return MSE + KLD