fix for dtypes

prishajain1 · prishajain1 · commit 516cf919b7b0 · 2025-12-27T10:39:09.000+05:30
diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
@@ -172,7 +172,7 @@ def __call__(
         padding_to_apply = tuple(current_padding)
         if any(p > 0 for dim_pads in padding_to_apply for p in dim_pads):
             x_padded = jnp.pad(
-                x_input, padding_to_apply, mode="constant", constant_values=0.0
+                x_input, padding_to_apply, mode="constant", constant_values=jnp.array(0.0, dtype=self.dtype)
             )
         else:
             x_padded = x_input

Original file line number	Diff line number	Diff line change
`@@ -172,7 +172,7 @@ def __call__(`
`172`	`172`	`padding_to_apply = tuple(current_padding)`
`173`	`173`	`if any(p > 0 for dim_pads in padding_to_apply for p in dim_pads):`
`174`	`174`	`x_padded = jnp.pad(`
`175`		`- x_input, padding_to_apply, mode="constant", constant_values=0.0`
	`175`	`+ x_input, padding_to_apply, mode="constant", constant_values=jnp.array(0.0, dtype=self.dtype)`
`176`	`176`	`)`
`177`	`177`	`else:`
`178`	`178`	`x_padded = x_input`