jpohhhh
/

biencoder_embedding

Feature Extraction

text-embeddings-inference

Model card Files Files and versions

jpohhhh commited on Apr 29, 2023

Commit

97be557

·

1 Parent(s): f2c3585

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -17,5 +17,5 @@ class EndpointHandler():
             A :obj:`list` | `dict`: will be serialized and returned
         """
         sentences = data.pop("inputs",data)
-        embeddings = self.model.encode(sentences, batch_size=100)
         return embeddings.tolist()

             A :obj:`list` | `dict`: will be serialized and returned
         """
         sentences = data.pop("inputs",data)
+        embeddings = self.model.encode(sentences, batch_size=100, device="cuda")
         return embeddings.tolist()