add model transforms

app.py

@@ -14,8 +14,11 @@ categories = ("Aculus Olearius", "Healthy", "Peacock Spot")
 def classify_health(input_img):
     input_img = transforms.ToTensor()(input_img)
     with torch.no_grad():
-        image = data_transforms(input_img).unsqueeze(0) 
-        probs = torch.nn.functional.softmax(model(image), dim=0)
+        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)
         idx = probs.argmax(dim=1)
     return dict(zip(categories, map(float, probs[0])))