Add load_best_model_at_end=True

bigcode-project · Jun 29, 2023 · d72c7fe · d72c7fe
1 parent 0d275cd
commit d72c7fe
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/finetune/finetune.py b/finetune/finetune.py
@@ -267,6 +267,8 @@ def run_training(args, train_data, val_data):
         output_dir=args.output_dir,
         dataloader_drop_last=True,
         evaluation_strategy="steps",
+        save_strategy="steps",
+        load_best_model_at_end=True,
         max_steps=args.max_steps,
         eval_steps=args.eval_freq,
         save_steps=args.save_freq,
@@ -309,4 +311,4 @@ def main(args):
 
     logging.set_verbosity_error()
 
-    main(args)
+    main(args)