diff --git a/django_htmx_refresh/middleware.py b/django_htmx_refresh/middleware.py
index 1ed3b84..fd10008 100644
--- a/django_htmx_refresh/middleware.py
+++ b/django_htmx_refresh/middleware.py
@@ -16,8 +16,8 @@ def process_template_response(self, request, response):
         if app_name in settings.HTMX_APPS:
             response.template_name = f"{app_name}/{response.template_name[0]}"
             if request.htmx:
-                response.context["base_template"] = "partial.html"
+                response.context_data["base_template"] = "partial.html"
             else:
-                response.context["base_template"] = "base.html"
+                response.context_data["base_template"] = "base.html"
         return response