debug for NaNs

prishajain1 · prishajain1 · commit 2fec0559e59e · 2025-12-29T20:59:16.000+05:30
diff --git a/src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py b/src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py
@@ -719,7 +719,7 @@ def step(
     )
 
     # Convert model_output (noise/v_pred) to x0_pred or epsilon_pred, based on prediction_type
-    model_output_for_history = self.convert_model_output(state, model_output, sample)
+    model_output_for_history = self.convert_model_output(state, model_output, sample, step_val)
     check_nan_jit(model_output_for_history, "model_output_for_history", step_val)
 
     # Apply corrector if applicable

Original file line number	Diff line number	Diff line change
`@@ -719,7 +719,7 @@ def step(`
`719`	`719`	`)`
`720`	`720`
`721`	`721`	`# Convert model_output (noise/v_pred) to x0_pred or epsilon_pred, based on prediction_type`
`722`		`- model_output_for_history = self.convert_model_output(state, model_output, sample)`
	`722`	`+ model_output_for_history = self.convert_model_output(state, model_output, sample, step_val)`
`723`	`723`	`check_nan_jit(model_output_for_history, "model_output_for_history", step_val)`
`724`	`724`
`725`	`725`	`# Apply corrector if applicable`