Function mean_flat

point_e/diffusion/gaussian_diffusion.py:1087–1091

Take the mean over all non-batch dimensions.

(tensor)

def mean_flat(tensor):
    """
    Take the mean over all non-batch dimensions.
    """
    return tensor.flatten(1).mean(1)