Fix cuda issue

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
2 years ago · 0e4870b6e8
1 changed files with 2 additions and 2 deletions
--- a/s_bert.py
+++ b/s_bert.py
@ -89,8 +89,8 @@ class STransformers(NNOperator):
        else:
            sentences = txt
        inputs = self.tokenize(sentences)
-        for k, v in inputs.items():
-            inputs[k] = v.to(self.device)
+#         for k, v in inputs.items():
+#             inputs[k] = v.to(self.device)
        embs = self.model(**inputs).cpu().detach().numpy()
        if isinstance(txt, str):
            embs = embs.squeeze(0)