Fix inputs

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
3 years ago · 92d43c7aa0
1 changed files with 1 additions and 1 deletions
--- a/codebert.py
+++ b/codebert.py
@ -77,7 +77,7 @@ class CodeBert(NNOperator):
            tokens = [self.tokenizer.cls_token, '<encoder-only>', self.tokenizer.sep_token] + tokens + \
                     [self.tokenizer.sep_token]
            tokens_ids = self.tokenizer.convert_tokens_to_ids(tokens)
-            inputs = torch.tensor(tokens_ids).to(self.device)
+            inputs = torch.tensor(tokens_ids).unsqueeze(0).to(self.device)
        except Exception as e:
            log.error(f'Invalid input for the tokenizer: {self.model_name}')
            raise e