Update app.py

app.py

@@ -175,7 +175,11 @@ def predict(image):
     try:
         image_tensor = preprocess_image(image)
         with torch.no_grad():
-            logits, * = model(image_tensor)
+            output = model(image_tensor)
+            if isinstance(output, tuple):
+                logits = output[0]
+            else:
+                logits = output
             _, predicted = torch.max(logits, 1)
         breed = dog_breeds[predicted.item()]  # Map label to breed name