Update app.py

app.py

@@ -40,11 +40,14 @@ class_labels = ['bird', 'cat', 'deer', 'dog', 'frog', 'horse']
 # Image transformation function
 def transform_image(image):
     """Preprocess the input image."""
+    mean, std = [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261]
+    img_size=224
     transform = transforms.Compose([
-        transforms.Resize(224),
+        transforms.Resize((img_size, img_size)),
         transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        transforms.Normalize(mean, std)
     ])
+
     img_tensor = transform(image).unsqueeze(0).to(device)
     return img_tensor