jpohhhh commited on
Commit
97be557
1 Parent(s): f2c3585

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -17,5 +17,5 @@ class EndpointHandler():
17
  A :obj:`list` | `dict`: will be serialized and returned
18
  """
19
  sentences = data.pop("inputs",data)
20
- embeddings = self.model.encode(sentences, batch_size=100)
21
  return embeddings.tolist()
 
17
  A :obj:`list` | `dict`: will be serialized and returned
18
  """
19
  sentences = data.pop("inputs",data)
20
+ embeddings = self.model.encode(sentences, batch_size=100, device="cuda")
21
  return embeddings.tolist()