diff --git a/celeryapp/tasks.py b/celeryapp/tasks.py
index 5e1690b31..76b2a4e71 100644
--- a/celeryapp/tasks.py
+++ b/celeryapp/tasks.py
@@ -16,6 +16,10 @@
 from daras_ai_v2.settings import templates
 from gooey_ui.pubsub import realtime_push
 from gooey_ui.state import set_query_params
+from gooeysite.bg_db_conn import db_middleware
+
+# workaround for - https://code.djangoproject.com/ticket/24810
+next_db_safe = db_middleware(next)
 
 
 @app.task
@@ -40,6 +44,7 @@ def gui_runner(
     error_msg = None
     set_query_params(query_params or {})
 
+    @db_middleware
     def save(done=False):
         if done:
             # clear run status
@@ -81,7 +86,7 @@ def save(done=False):
             start_time = time()
             try:
                 # advance the generator (to further progress of run())
-                yield_val = next(gen)
+                yield_val = next_db_safe(gen)
                 # increment total time taken after every iteration
                 run_time += time() - start_time
                 continue