add model transforms

app.py

@@ -17,7 +17,7 @@ def classify_health(input_img):
         image = data_transforms(input_img).unsqueeze(0) 
         probs = model(image)
         idx = probs.argmax(dim=1)
-    return dict(zip(categories, map(float, probs)))
+    return dict(zip(categories, map(float, probs[0])))
 
 
 labels = gr.Label()