Update app.py

app.py

@@ -24,7 +24,7 @@ def predict(inputs: str) -> Dict[str, Any]:
     probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
 
     # Get probabilities for each label
-    label_probabilities = {label: round(prob.item(), 4) for label, prob in zip(labels, probabilities[0].tolist())}
+    label_probabilities = {label: round(prob, 4) for label, prob in zip(labels, probabilities[0].tolist())}
 
     # Return the result
     return label_probabilities