def normalize(img, radius=1):
    """
    Normalize image to L2 norm of radius
    """
    img = img - img.mean(-2, keepdim=True).mean(-1, keepdim=True)
    norms = img.pow(2).sum(-2, keepdim=True).sum(-1, keepdim=True).sqrt()
    return img / norms * radius