From 5adce7aad0da6949c59723ffc0754802390b63a6 Mon Sep 17 00:00:00 2001
From: Jack Roberts <jroberts@turing.ac.uk>
Date: Thu, 25 Jul 2024 20:57:05 +0100
Subject: [PATCH] set max_new_tokens for generate

---
 scripts/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scripts/train.py b/scripts/train.py
index 2e82baae..80d9e953 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -105,6 +105,7 @@ def main(experiment_path):
             ],
             accelerator=Accelerator(),
             n_print=5,
+            max_new_tokens=50,
         )
 
     # Step 7: Load trainer