import torch


def loss_kl(z_mean, z_logvar):
    """
    Args:
      z_mean of shape (batch_size, n_components): Means of the approximate distributions of the codes.
      z_logvar of shape (batch_size, n_components): Log-variance of the approximate distributions of the codes.
    
    Returns:
      loss (torch scalar): Kullback-Leibler divergence.
    """
    loss = (0.5 * torch.sum(torch.exp(z_logvar) + z_mean**2 - 1.0 - z_logvar))/z_mean.shape[0]
    return loss