fix profiling

coolkp · coolkp · commit 1203400f193a · 2025-06-12T07:01:17.000Z
diff --git a/src/maxdiffusion/trainers/sdxl_trainer.py b/src/maxdiffusion/trainers/sdxl_trainer.py
@@ -223,6 +223,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
         (unet_state, train_metric, train_rngs) = p_train_step(
             unet_state, vae_state, text_encoder_state, text_encoder_2_state, example_batch, train_rngs
         )
+        train_metric['scalar']['learning/loss'].block_until_ready()
 
       samples_count = self.total_train_batch_size * (step + 1)
       new_time = datetime.datetime.now()

Original file line number	Diff line number	Diff line change
`@@ -223,6 +223,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera`
`223`	`223`	`(unet_state, train_metric, train_rngs) = p_train_step(`
`224`	`224`	`unet_state, vae_state, text_encoder_state, text_encoder_2_state, example_batch, train_rngs`
`225`	`225`	`)`
	`226`	`+ train_metric['scalar']['learning/loss'].block_until_ready()`
`226`	`227`
`227`	`228`	`samples_count = self.total_train_batch_size * (step + 1)`
`228`	`229`	`new_time = datetime.datetime.now()`