Fix W&B validation step logging

algorithms/dememwm/df_video.py

@@ -1183,7 +1183,7 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                     log_video(
                         xs_pred_decode,
                         xs_decode,
-                        step=batch_idx,
+                        step=getattr(self, "global_step", 0),
                         namespace=namespace + "_vis",
                         context_frames=self.context_frames,
                         logger=self.logger.experiment,

main.py

@@ -224,6 +224,9 @@ def run_local(cfg: DictConfig):
             id=wandb_run_id,
             resume="auto"
         )
+        if is_rank_zero:
+            logger.experiment.define_metric("trainer/global_step")
+            logger.experiment.define_metric("*", step_metric="trainer/global_step")
 
     else:
         logger = None