Update app.py

app.py

@@ -39,7 +39,7 @@ def inference(input_image):
     with torch.no_grad():
         output = model(input_batch)
     # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
-    probabilities = torch.nn.functional.softmax(output[0], dim=0)
+    probabilities = torch.nn.functional.softmax(output[0])
 
     # Read the categories
     #with open("dog_cat.txt", "r") as f: