Update app.py

app.py

@@ -67,7 +67,7 @@ def main():
             logits = model(input_ids=input_ids, attention_mask=attention_mask)
             _, prediction = torch.max(logits, 1)
             prediction = prediction.item()
-            predict_proba = round(torch.softmax(logits, 1).numpy().squeeze()[prediction].item(),4)
+            predict_proba = round(torch.softmax(logits, 1).cpu().squeeze().tolist()[prediction],4)
             return prediction, predict_proba
 
     def pred_str(prediction):