Browse Source
Fix cuda issue
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
2 additions and
2 deletions
-
s_bert.py
|
|
@ -89,8 +89,8 @@ class STransformers(NNOperator): |
|
|
else: |
|
|
else: |
|
|
sentences = txt |
|
|
sentences = txt |
|
|
inputs = self.tokenize(sentences) |
|
|
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() |
|
|
embs = self.model(**inputs).cpu().detach().numpy() |
|
|
if isinstance(txt, str): |
|
|
if isinstance(txt, str): |
|
|
embs = embs.squeeze(0) |
|
|
embs = embs.squeeze(0) |
|
|
|
