[tokenizer] Uses fp32 for TextembeddingTranslator clip() (#2881)

deepjavalibrary · Dec 1, 2023 · 2fc8058 · 2fc8058
1 parent c057898
commit 2fc8058
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/...sions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/...sions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -113,7 +113,7 @@ private static NDArray meanPool(NDArray embeddings, NDArray attentionMask, boole
         long[] shape = embeddings.getShape().getShape();
         attentionMask = attentionMask.expandDims(-1).broadcast(shape);
         NDArray inputAttentionMaskSum = attentionMask.sum(AXIS);
-        NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12);
+        NDArray clamp = inputAttentionMaskSum.clip(1e-9f, 1e12f);
         NDArray prod = embeddings.mul(attentionMask);
         NDArray sum = prod.sum(AXIS);
         if (sqrt) {