import torch


def _kl_loss(bn_mean, bn_std, tmp_mean, tmp_std):
    """
    Compute KL divergence given two Gaussian distributions whose statistics
    are (bn_meanm, bn_std) and (tmp_mean, tmp_std), respectively.
    """
    a = torch.log(tmp_std / bn_std)
    c = (bn_std ** 2 + (bn_mean - tmp_mean) ** 2) / tmp_std ** 2
    b = 0.5 * (1 - c)
    loss = a - b
    return loss.mean()