added nnx.jit decorator to encode and decode

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -1193,6 +1193,7 @@ def __init__(
         precision=precision,
     )
 
+  @nnx.jit
   def encode(
       self, x: jax.Array, return_dict: bool = True
   ) -> Union[FlaxAutoencoderKLOutput, Tuple[FlaxDiagonalGaussianDistribution]]:
@@ -1225,6 +1226,7 @@ def scan_fn(carry, input_slice):
         return (posterior,)
     return FlaxAutoencoderKLOutput(latent_dist=posterior)
 
+  @nnx.jit
   def decode(
       self, z: jax.Array, return_dict: bool = True
   ) -> Union[FlaxDecoderOutput, jax.Array]: