diff --git a/ee/api/chalicelib/core/heatmaps.py b/ee/api/chalicelib/core/heatmaps.py
index 68ced64f1..c6e7c3169 100644
--- a/ee/api/chalicelib/core/heatmaps.py
+++ b/ee/api/chalicelib/core/heatmaps.py
@@ -90,7 +90,7 @@ def get_by_url(project_id, data: schemas.GetHeatmapPayloadSchema):
         return helper.list_to_camel_case(rows)
 
 
-def get_x_y_by_url_and_session_id(project_id, session_id, data: schemas.GetHeatmapBasePayloadSchema):
+def get_x_y_by_url_and_session_id(project_id, session_id, data: schemas.GetHeatmapPayloadSchema):
     args = {"project_id": project_id, "session_id": session_id, "url": data.url}
     constraints = ["main_events.project_id = toUInt16(%(project_id)s)",
                    "main_events.session_id = %(session_id)s",
@@ -120,7 +120,7 @@ def get_x_y_by_url_and_session_id(project_id, session_id, data: schemas.GetHeatm
         return helper.list_to_camel_case(rows)
 
 
-def get_selectors_by_url_and_session_id(project_id, session_id, data: schemas.GetHeatmapBasePayloadSchema):
+def get_selectors_by_url_and_session_id(project_id, session_id, data: schemas.GetHeatmapPayloadSchema):
     args = {"project_id": project_id, "session_id": session_id, "url": data.url}
     constraints = ["main_events.project_id = toUInt16(%(project_id)s)",
                    "main_events.session_id = %(session_id)s",