add eval time

Author: susanbao

src/maxdiffusion/trainers/wan_trainer.py

@@ -328,6 +328,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
           while True:
             try:
               with mesh:
+                eval_start_time = datetime.datetime.now()
                 eval_batch = load_next_batch(eval_data_iterator, None, self.config)
                 metrics, eval_rng = p_eval_step(state, eval_batch, eval_rng, scheduler_state)
                 losses = metrics["scalar"]["learning/eval_loss"]
@@ -336,11 +337,15 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
                 gathered_timesteps = jax.device_get(gathered_timesteps_on_device)
                 gathered_losses_on_device = multihost_utils.process_allgather(losses)
                 gathered_losses = jax.device_get(gathered_losses_on_device)
-                for t, l in zip(gathered_timesteps, gathered_losses):
+                for t, l in zip(gathered_timesteps.flatten(), gathered_losses.flatten()):
                   timestep = int(t)
                   if timestep not in eval_losses_by_timestep:
                       eval_losses_by_timestep[timestep] = []
                   eval_losses_by_timestep[timestep].append(l)
+                eval_end_time = datetime.datetime.now()
+                eval_duration = eval_end_time - eval_start_time
+                if jax.process_index() == 0:
+                  max_logging.log(f"  Eval step time {eval_duration.total_seconds():.2f} seconds.")
             except StopIteration:
               # This block is executed when the iterator has no more data
               break