fix the context manager call

2023-09-03 22:49:28 -04:00
parent 5e4a760ad8
commit 83d904a27d
1 changed files with 1 additions and 1 deletions
--- a/src/axolotl/utils/callbacks.py
+++ b/src/axolotl/utils/callbacks.py
@@ -210,7 +210,7 @@ def bench_eval_callback_factory(trainer, tokenizer):
            "subject": example["subject"],
        }

-    with dist_state.main_process_first:
+    with dist_state.main_process_first():
        bench_dataset = bench_dataset.map(tokenize_evals)
        bench_dataset = bench_dataset.filter(lambda x: x["labels"][-2] in abcd_idx)