Update app.py

app.py

@@ -104,7 +104,7 @@ def ant_genus_classification(image):
         # preds = torch.argmax(preds, dim=1)
         print(preds.shape)
         
-        confidences = {class_to_idx[i]: float(preds[i]) for i in range(len(classes))}
+        confidences = {class_to_idx[i]: float(preds[0][i]) for i in range(len(classes))}
     return confidences
     # prediction = model(image)[0]
     # prediction = prediction.tolist()