Update inference.py

inference.py

@@ -28,7 +28,9 @@ def inference(input_text):
     tokenizer = torch.load(saved_model_path + 'fine-tuned-tokenizer.pt', map_location=torch.device('cpu'))
     model.eval()
     input = tokenizer(input_text)
-    output = model(input['input_ids'], input['attention_mask'])
+    input_ids = torch.Tensor(input['input_ids']).squeeze(1)
+    mask = torch.Tensor(input['attention_mask'])
+    output = model(input_ids, mask)
     
     return output.tolist()