Cast embeddings to float32 before computing distances.

This fixes a bug where bfloat16 is not supported in nonGPU or nonTPU machine. PiperOrigin-RevId: 646651355
PAIR-code · Jun 25, 2024 · 5456011 · 5456011
1 parent 2e9d267
commit 5456011
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/lit_nlp/components/nearest_neighbors.py b/lit_nlp/components/nearest_neighbors.py
@@ -124,7 +124,11 @@ def run(
 
     # <float32>[emb_size]
     dataset_embs = [output[nnconf.embedding_name] for output in dataset_outputs]
+    dataset_embs = [emb.astype(np.float32) for emb in dataset_embs]
+
     example_embs = [example_output[nnconf.embedding_name]]
+    example_embs = [emb.astype(np.float32) for emb in example_embs]
+
     distances = distance.cdist(example_embs, dataset_embs)[0]
     sorted_indices = np.argsort(distances)
     k = nnconf.num_neighbors