diff --git a/modeling/inference_model.py b/modeling/inference_model.py
index 1bf35146..5b40733e 100644
--- a/modeling/inference_model.py
+++ b/modeling/inference_model.py
@@ -365,8 +365,8 @@ class InferenceModel:
                     # stop temporarily to insert WI, we can assume that we are done
                     # generating. We shall break.
                     if (
-                        self.gen_state["halt"]
-                        or not self.gen_state["regeneration_required"]
+                        self.gen_state.get("halt")
+                        or not self.gen_state.get("regeneration_required")
                     ):
                         break