Update app.py

app.py

@@ -55,8 +55,8 @@ def inference(title, abstract, threshold=0.95):
     attention_mask = encoding["attention_mask"].to(device)
 
     with torch.no_grad():
-        res_probs = class_model(input_ids, attention_mask) # shape: (1, 8)
-
+        res_probs = torch.exp(class_model(input_ids, attention_mask))
+    
     probs = res_probs.squeeze(0)  # (8,)
     sorted_probs, sorted_indices = torch.sort(probs, descending=True)