SentenceTransformer encode with torch

snap-stanford · Aug 27, 2024 · 89c91de · 89c91de
1 parent 56fe85e
commit 89c91de
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/text_embedder.py b/examples/text_embedder.py
@@ -15,4 +15,4 @@ def __init__(self, device: Optional[torch.device] = None):
         )
 
     def __call__(self, sentences: List[str]) -> Tensor:
-        return torch.from_numpy(self.model.encode(sentences))
+        return self.model.encode(sentences, convert_to_tensor=True)