From fa0786f6b55f2049c8fe2a70ba20e439f68f0561 Mon Sep 17 00:00:00 2001 From: David Hall Date: Tue, 26 Nov 2024 09:08:18 -0800 Subject: [PATCH] comment out GradWatchCallback --- src/levanter/main/train_lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/levanter/main/train_lm.py b/src/levanter/main/train_lm.py index 3fb1c6afa..a827ae46b 100644 --- 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))