def split_mean_var(output, mean_act, k):
    """
    Format NN output into mean and variances (or in case of Laplace llh: loc. parameter, 2 * scale paramaeter)
    :param output: (batch, k * 2 tensor:
                    k elements: mean, then k elements: variances
    :param mean_act: activation function that will be applied to the mean
    :param k: dimension of regression output
    :return: mean, log variances
    """
    assert k > 0, "k must be positive!"
    assert output.shape[1] == k * 2, "Aleatoric uncertainty estimation: wrong input shape!"
    mean = mean_act(output[:, :k])
    logvar = output[:, k:]
    return mean, logvar