fix during training eval config

jbloomAus · Oct 18, 2024 · 333d71c · 333d71c
1 parent 9341398
commit 333d71c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/sae_lens/training/sae_trainer.py b/sae_lens/training/sae_trainer.py
@@ -136,9 +136,12 @@ def __init__(
         self.trainer_eval_config = EvalConfig(
             batch_size_prompts=self.cfg.eval_batch_size_prompts,
             n_eval_reconstruction_batches=self.cfg.n_eval_batches,
+            n_eval_sparsity_variance_batches=self.cfg.n_eval_batches,
             compute_ce_loss=True,
-            n_eval_sparsity_variance_batches=1,
             compute_l2_norms=True,
+            compute_sparsity_metrics=True,
+            compute_variance_metrics=True,
+            compute_kl=False,
         )
 
     @property