import torch


def kl(z1_log_var, z1_mean):
    """kl loss for VAE

        Parameters
        ----------
        z1_log_var: torch.Tensor
            Tensor of shape [batch_size, z_dim]
        z1_mean: torch.Tensor
            Tensor of shape [batch_size, z_dim]

        Returns
        -------
        Returns the computed kl loss.
    """
    z1_var = torch.exp(z1_log_var) + 1e-4
    return -0.5 * torch.sum(1 + z1_log_var - z1_mean ** 2 - z1_var)