fix for dtypes

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -274,6 +274,7 @@ def __init__(
         weights_dtype: jnp.dtype = jnp.float32,
         precision: jax.lax.Precision = None,
     ):
+        self.dtype = dtype
         self.dim = dim
         self.mode = mode