comment out GradWatchCallback

stanford-crfm · Nov 26, 2024 · fa0786f · fa0786f
1 parent 0a7d78f
commit fa0786f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/main/train_lm.py b/src/levanter/main/train_lm.py
@@ -170,7 +170,7 @@ def main(config: TrainLmConfig):
                 every=config.hf_save_steps,
             )
 
-        trainer.add_hook(callbacks.GradWatchCallback(include_histogram=True), every=5)
+        # trainer.add_hook(callbacks.GradWatchCallback(include_histogram=True), every=5)
 
         state = trainer.initial_state(training_key, model_init=lambda: config.model.build(Vocab, key=model_key))