Fix

app.py

@@ -47,20 +47,17 @@ model.to(device)
 def predict(image):
     try:
         transform = transforms.Compose([
-            transforms.Resize(256),
-            transforms.CenterCrop(224),
-            transforms.ToTensor(),  # Converts to [0, 1] range
-            transforms.Normalize(
-                mean=[0.485, 0.456, 0.406],
-                std=[0.229, 0.224, 0.225]
-            )
+                    transforms.Resize((224, 224)),  # match ViT input size
+                        transforms.Normalize(mean=(0.485, 0.456, 0.406),
+                             std=(0.229, 0.224, 0.225))
         ])
 
         pixel_values = transform(image).unsqueeze(0).to(device)
 
         with torch.no_grad():
             logits = model(pixel_values).logits
-            probs = torch.nn.functional.softmax(logits, dim=-1).squeeze()
+            probs = torch.nn.functional.softmax(logits, dim=1).squeeze()
+
 
         topk = torch.topk(probs, k=5)
         top_indices = topk.indices.tolist()