diff --git a/scripts/base_train.py b/scripts/base_train.py
index b691ed4..518bd71 100644
--- a/scripts/base_train.py
+++ b/scripts/base_train.py
@@ -219,7 +219,7 @@ for step in range(num_iterations + 1):
             "My favorite color is",
             "If 5*x + 3 = 13, then x is",
         ]
-        engine = Engine(model, tokenizer)
+        engine = Engine(orig_model, tokenizer)
         for prompt in prompts:
             tokens = tokenizer(prompt, prepend="<|bos|>")
             with autocast_ctx: