Added image

app.py

@@ -35,7 +35,7 @@ def predict(img):
     preds = torch.softmax(pred_logits, dim=1)
     
     # Create a prediction label and prediction probability dictionary for each prediction class
-    pred_and_prob_labels = {class_names[i]: preds[i].item() for i in range(len(class_names))}
+    pred_and_prob_labels = {class_names[i]: preds[0][i].item() for i in range(len(class_names))}
 
   return pred_and_prob_labels