Add missing batch size argument

adamkarvonen · Dec 4, 2024 · 877f2e7 · 877f2e7
1 parent e0cb629
commit 877f2e7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/custom_saes/run_all_evals_custom_saes.py b/custom_saes/run_all_evals_custom_saes.py
@@ -146,7 +146,10 @@ def run_evals(
         "unlearning": (
             lambda: unlearning.run_eval(
                 unlearning.UnlearningEvalConfig(
-                    model_name="gemma-2-2b-it", random_seed=RANDOM_SEED, llm_dtype=llm_dtype
+                    model_name="gemma-2-2b-it",
+                    random_seed=RANDOM_SEED,
+                    llm_dtype=llm_dtype,
+                    llm_batch_size=llm_batch_size,
                 ),
                 selected_saes,
                 device,