Update inference.py

inference.py

@@ -31,7 +31,8 @@ def inference(input_text):
     input_ids = torch.Tensor(input['input_ids']).to(torch.device('cpu')).long()
     input_ids.resize_(1,len(input_ids))
     print(input_ids)
-    mask = torch.Tensor(input['attention_mask'])
+    mask = torch.Tensor(input['attention_mask']).to(torch.device('cpu'))
+    mask.resize_(1, len(mask))
     output = model(input_ids, mask)
     
     return output.tolist()