add model transforms

app.py

@@ -16,9 +16,7 @@ def classify_health(input_img):
     with torch.no_grad():
         image = data_transforms(input_img).unsqueeze(0)
         output = model(image)
-        print(output)
-        print(output.shape)
-        probs = torch.nn.functional.softmax(model(image)[0], dim=0)
+        probs = torch.nn.functional.softmax(output, dim=1)
         idx = probs.argmax(dim=1)
     return dict(zip(categories, map(float, probs[0])))