diff --git a/finetune.py b/finetune.py
index 4059fc7..ba019ff 100644
--- a/finetune.py
+++ b/finetune.py
@@ -123,7 +123,7 @@ def train(
     )
     tokenizer.padding_side = "left"  # Allow batched inference
 
-    def tokenize(prompt):
+    def tokenize(prompt, add_eos_token=True):
         # there's probably a way to do this with the tokenizer settings
         # but again, gotta move fast
         result = tokenizer(