set max_new_tokens for generate

alan-turing-institute · Jul 25, 2024 · 5adce7a · 5adce7a
1 parent 3c90a06
commit 5adce7a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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