sigh

stanford-crfm · Feb 13, 2024 · 2cddfd5 · 2cddfd5
1 parent c241da4
commit 2cddfd5
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/levanter/trainer.py b/src/levanter/trainer.py
@@ -358,7 +358,6 @@ def init_state(partial_state, model_init, training_key, is_trainable):
             return eqx.combine(partial_state, fresh_state)
 
         state = init_state(state, model_init, training_key, self.is_trainable_param)
-        print(state.step.sharding)
 
         return state
 
@@ -385,16 +384,16 @@ def training_steps(
             with capture_time() as loading_time:
                 example = next(iter_data)
 
-            levanter.tracker.log_metrics({"throughput/loading_time": loading_time()}, step=state.step)
-
             info = self.train_step(state, example)
             state = info.state
 
             if run_hooks:
                 with capture_time() as hook_time:
                     self.run_hooks(info)
 
-                levanter.tracker.log_metrics({"throughput/hook_time": hook_time()}, step=state.step)
+                levanter.tracker.log_metrics({"throughput/hook_time": hook_time()}, step=info.step)
+
+            levanter.tracker.log_metrics({"throughput/loading_time": loading_time()}, step=info.step)
 
             yield info