def gaussian_neg_log_likelihood_diag(parameters, target):
    """Negative log likelihood loss for the Gaussian distribution
    B = batch size, D = dimension of target (num classes), N = ensemble size

    Args:
        parameters (torch.tensor((B, D)), torch.tensor((B, D))):
            mean values and variances of y|x for every x in
            batch.

        target (torch.tensor((B, N, D))): sample from the normal
            distribution, if not an ensemble prediction N=1.
    """
    B, N, _ = target.size()
    mu, sigma_sq = parameters

    prec = sigma_sq.pow(-1)

    nll = 0.0
    for mu_b, prec_b, target_b in zip(mu, prec, target):
        sample_var = (target_b - mu_b).var(dim=0)
        trace_term = (prec_b * sample_var).sum() * N / 2
        nll += trace_term - N / 2 * prec_b.prod()

    return nll / B