update stuff

app.py

@@ -31,9 +31,11 @@ state_dict = torch.load('model.pth', map_location='cpu')
 model.load_state_dict(state_dict, strict=False)
 model.eval()
 
-def predict(im):
-    im = im['background']
-    print(im)
+def predict(image):
+    image = image['composite']
+    image = transforms.Resize((28, 28))(image)
+    print(image)
+    
     x = torch.tensor(im, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.
 
     with torch.no_grad():