Update app.py

app.py

@@ -28,7 +28,7 @@ def forward_pass(title, abstract, tokenizer, model):
     with torch.no_grad():
         logits = model(embeddings[None])['logits'][0]
         assert logits.shape == (8,)
-        probs = torch.softmax(logits).data.cpu().numpy()
+        probs = torch.softmax(logits, dim=0).data.cpu().numpy()
     
     return probs