From 6dafe84fda72c5f34cf1f3e608394f1f87a05742 Mon Sep 17 00:00:00 2001 From: Ddupg Date: Thu, 2 Jul 2026 21:58:46 +0800 Subject: [PATCH] fix(python): avoid CPU bf16 compile flake --- python/python/lance/torch/distance.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/python/lance/torch/distance.py b/python/python/lance/torch/distance.py index 81201027c87..ed5292cdb2f 100644 --- a/python/python/lance/torch/distance.py +++ b/python/python/lance/torch/distance.py @@ -150,8 +150,9 @@ def pairwise_l2( if x.dtype != y.dtype or (y2 is not None and x.dtype != y2.dtype): raise ValueError("pairwise_l2 data types do not match") origin_dtype = x.dtype - if x.device == torch.device("cpu") and x.dtype == torch.float16: - # Pytorch does not support `x @ y.T` for float16 on CPU + if x.device == torch.device("cpu") and x.dtype in (torch.float16, torch.bfloat16): + # Pytorch does not support `x @ y.T` for float16 on CPU, and bfloat16 + # can hit backend-specific compiler paths. Accumulate in float32. x = x.type(torch.float32) y = y.type(torch.float32) if y2 is not None: