Update app.py

app.py

@@ -152,7 +152,7 @@ def predict(image):
                 logits = output[0]
             else:
                 logits = output
-            _, predicted = torch.max(logits, 1) # predicted is the max value's index
+            _, predicted = torch.max(logits, dim=1) # predicted is the max value's index
         breed = dog_breeds[predicted.item()]
         
         description = get_dog_description(breed)