don't change loop while running

dbanys · May 7, 2021 · c648e20 · c648e20
1 parent d13a182
commit c648e20
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/aitextgen/aitextgen.py b/aitextgen/aitextgen.py
@@ -403,9 +403,10 @@ def generate(
 
                     # remove fields not in schema_return
                     if schema_return:
+                        keys = gen_text_dict.keys()
                         if len(schema_return) == 1:
                             gen_text_dict = gen_text_dict[schema_return[0]]
-                        for key in gen_text_dict.keys():
+                        for key in keys:
                             if key not in schema_return:
                                 gen_text_dict.pop(key, None)