1 files changed, 2 insertions, 2 deletions
diff --git a/ai_dungeon.py b/ai_dungeon.py
index 2b976f6..baac638 100755
--- a/ai_dungeon.py
+++ b/ai_dungeon.py
@@ -110,12 +110,12 @@ def main() -> None:
     args = parser.parse_args()
     model = load_model(args.model)
     script = pick_flavor()
-    prologue = generate(
+    prologue = wrap(generate(
         model,
         script,
         max_new_tokens=20,
         forced_eos_token_id=model.tokenizer.eos_token_id,
-    )
+    ))
     print(prologue)
     prompt = prologue
     prev_response_length = 0