debug for NaNs

prishajain1 · prishajain1 · commit 3d5957de1a51 · 2025-12-29T21:03:41.000+05:30
diff --git a/src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py b/src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py
@@ -731,6 +731,7 @@ def step(
             last_sample=state.last_sample,
             this_sample=sample,
             order=state.this_order,
+            step=step_val
         ),
         lambda: sample,
     )

Original file line number	Diff line number	Diff line change
`@@ -731,6 +731,7 @@ def step(`
`731`	`731`	`last_sample=state.last_sample,`
`732`	`732`	`this_sample=sample,`
`733`	`733`	`order=state.this_order,`
	`734`	`+ step=step_val`
`734`	`735`	`),`
`735`	`736`	`lambda: sample,`
`736`	`737`	`)`