fix for dtypes

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -511,7 +511,10 @@ def __call__(self, x: jax.Array, cache: Dict[str, Any] = None):
         input_dtype = x.dtype
 
         h, sc_cache = self.conv_shortcut(x, cache.get("shortcut"))
-        new_cache["shortcut"] = sc_cache.astype(self.dtype)
+        if sc_cache is not None:
+            new_cache["shortcut"] = sc_cache.astype(self.dtype)
+        else:
+            new_cache["shortcut"] = None
 
         x = self.norm1(x)
         x = self.nonlinearity(x)