[tokenizer] Fixes TextEmbeddingTranslator only uses CPU (#3122)

deepjavalibrary · Apr 24, 2024 · 6b7215c · 6b7215c
1 parent a2f075f
commit 6b7215c
Showing 1 changed file with 0 additions 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
@@ -175,7 +175,6 @@ private NDArray processEmbedding(NDManager manager, NDList list, Encoding encodi
                 throw new AssertionError("Unexpected pooling mode: " + pooling);
         }
 
-        embedding = embedding.toDevice(Device.cpu(), false);
         if (denseModel != null) {
             NDArray weight = denseModel.get("linear.weight");
             NDArray bias = denseModel.get("linear.bias");