Update inference.py
Browse files- inference.py +1 -1
inference.py
CHANGED
|
@@ -28,7 +28,7 @@ def inference(input_text):
|
|
| 28 |
tokenizer = torch.load(saved_model_path + 'fine-tuned-tokenizer.pt', map_location=torch.device('cpu'))
|
| 29 |
model.eval()
|
| 30 |
input = tokenizer(input_text)
|
| 31 |
-
output = model(
|
| 32 |
|
| 33 |
return output.tolist()
|
| 34 |
|
|
|
|
| 28 |
tokenizer = torch.load(saved_model_path + 'fine-tuned-tokenizer.pt', map_location=torch.device('cpu'))
|
| 29 |
model.eval()
|
| 30 |
input = tokenizer(input_text)
|
| 31 |
+
output = model(input['input_ids'].squeeze(1), input['attention_mask'])
|
| 32 |
|
| 33 |
return output.tolist()
|
| 34 |
|