kdevoe commited on
Commit
da71a7f
·
verified ·
1 Parent(s): a0eab2d

Update inference.py

Browse files
Files changed (1) hide show
  1. 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(input['input_ids'].squeeze(1), input['attention_mask'])
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'], input['attention_mask'])
32
 
33
  return output.tolist()
34