diff --git a/s_bert.py b/s_bert.py index 1a73b58..7586c2e 100644 --- a/s_bert.py +++ b/s_bert.py @@ -85,8 +85,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)