def torch_mean_sum(x, n_mean=1):
    """ Average over batch dim, sum over all other dims. """
    sum_dims = tuple(range(n_mean, x.ndim))

    if sum_dims:
        x = x.sum(sum_dims)

    return x.mean()