fix bug if num_samples<100

ai-hero · May 26, 2024 · b336061 · b336061
1 parent 30f3ce8
commit b336061
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/aihero/research/finetuning/infer.py b/src/aihero/research/finetuning/infer.py
@@ -16,8 +16,7 @@
 
 CHECKPOINT_DIR = "/mnt/checkpoint"
 DATASET_DIR = "/mnt/dataset"
-MAX_NEW_TOKENS = 512
-
+MAX_NEW_TOKENS = 7680
 if os.environ.get("HF_TOKEN", None):
     login(token=os.environ["HF_TOKEN"])