debug for NaNs

prishajain1 · prishajain1 · commit 9474c154759f · 2025-12-29T21:06:40.000+05:30
diff --git a/src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py b/src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py
@@ -786,6 +786,7 @@ def non_step_idx0_branch():
         model_output=model_output,
         sample=sample,
         order=state.this_order,
+        step=step_val,
     )
 
     # Update lower_order_nums for warmup

Original file line number	Diff line number	Diff line change
`@@ -786,6 +786,7 @@ def non_step_idx0_branch():`
`786`	`786`	`model_output=model_output,`
`787`	`787`	`sample=sample,`
`788`	`788`	`order=state.this_order,`
	`789`	`+ step=step_val,`
`789`	`790`	`)`
`790`	`791`
`791`	`792`	`# Update lower_order_nums for warmup`