Update app.py

app.py

@@ -46,7 +46,7 @@ result
 # transform logit to get probabilities
 if (result.logits.size(dim=1) < 2):
     pad = (0, 1)
-    result = nn.functional.pad(result, pad, "constant", 0)
+    result = nn.functional.pad(result.logits, pad, "constant", 0)
     st.write(result)
     prediction = nn.functional.softmax(result, dim=-1)
 prediction = nn.functional.softmax(result.logits, dim=-1)