Update app.py

app.py

@@ -20,10 +20,10 @@ labels = r.text.split("\n")
 
 # Define prediction function
 def predict(img):
-    # Transform image to pytorch tensor of shape [1, 224, 224, 3]
-    img = T.ToTensor()(img).unsqueeze(0)
-    img = T.Resize(size=(224, 224))
-    print(img.shape)
+    # Transform image to pytorch tensor of shape [1, 3, 224, 224]
+    img = T.PILToTensor()(img).unsqueeze(0)
+    img = T.Resize(size=(224, 224))(img)
+
     
     # Use model without gradients to reduce computation
     with torch.no_grad():