import torch


def kl_divergence(mu, logvar):
    """
    Returns the KLD of the sampled latent vector w.r.t. to a Standard Normal Gaussian

    INPUT:
    mu: Tensor (B, 1) -> mean of a Gaussian distribution q(z|x_test)
    logvar: Tensor (B, 1) -> logvar of a Gaussian distribution q(z|x_test)

    OUTPUT:
    kld_loss: Tensor(B, 1) -> Kullback-Leibler divergence
    """

    kld_loss = torch.mean(-0.5 * torch.sum(1 + logvar - mu ** 2 - logvar.exp(), dim=1), dim=0)

    return kld_loss