Update app.py

app.py

@@ -30,8 +30,8 @@ def predict(image):
     
     # Perform post-processing if needed (e.g., softmax for probabilities)
     # Replace this with your actual post-processing logic
-    print(output.logits.argmax(1).item())
-    probabilities = torch.softmax(output, dim=1).squeeze().tolist()
+    probabilities = torch.softmax(output.logits, dim=1).squeeze().tolist()
+
     
     # Map the class indices to class labels
     class_labels = ["Class1", "Class2", "Class3", "Class4"]