def euclidean_distance_nosqrt(x1, x2):
    """ 欧氏距离，不对结果开方

    Args:
        x1: [B, N] or [N]
        x2: same shape as x1

    Returns:
        [B] vector or scalar
    """
    return (x1 - x2).pow(2).sum(-1)