Update app.py

app.py

@@ -25,7 +25,8 @@ optimizer = torch.optim.Adam(params = TinyCNN_model.parameters() ,lr=0.001)
 
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
-    transforms.Grayscale(num_output_channels=3)
+    transforms.Grayscale(num_output_channels=3),
+    transforms.ToTensor()
 ])
 
 # Load saved weights
@@ -39,7 +40,7 @@ TinyCNN_model.load_state_dict(
 ### 3. Predict function ###
 
 # Create predict function
-def predict(img) -> Tuple[Dict, float]:
+def predict(img) :
     """Transforms and performs a prediction on img and returns prediction and time taken.
     """
     # Start the timer