diff --git a/deployment/handlers/prediction_app.py b/deployment/handlers/prediction_app.py
index 7deae3a..7555934 100644
--- a/deployment/handlers/prediction_app.py
+++ b/deployment/handlers/prediction_app.py
@@ -90,7 +90,9 @@ def handler(event, context):
             return {
                 "statusCode": 200,
                 "body": json.dumps(
-                    {"message": f"{model.predict(data=pd.DataFrame({'a': [1]}))}"}
+                    {
+                        "message": f"{model.predict(data=pd.DataFrame(columns=model.model.original_features))}"
+                    }
                 ),
             }
         else: