Update app.py

app.py

@@ -80,7 +80,7 @@ def image_transformation(image):
 
 def image_prediction(image, model):
     image = image_transformation(image=image)
-    outputs = model.predict(image)
+    outputs = double(model.predict(image))
     _, y_hat = outputs.max(1)
     predicted_idx = str(y_hat.item())
     return predicted_idx
@@ -106,7 +106,7 @@ def main():
 
         if pred_button:
             image_prediction(image=image, model=model)
-            outputs = int(model.predict(image_prediction))
+            outputs = double(model.predict(image_prediction))
             _, y_hat = outputs.max(1)
             predicted_idx = str(y_hat.item())
             right_column.title("Prediction")