Update app.py

app.py

@@ -42,7 +42,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], dim=1)
 
     # Read the categories
     with open("artist_classes.txt", "r") as f: