From 92d43c7aa0918bca6bc01fcd91101879007f90b5 Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Tue, 8 Nov 2022 16:57:28 +0800
Subject: [PATCH] Fix inputs

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
---
 codebert.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/codebert.py b/codebert.py
index 274e1a2..c099494 100644
--- 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