removed nnx.jit for debug

prishajain1 · prishajain1 · commit 4cac7996235b · 2025-12-28T11:17:17.000+05:30
diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
@@ -1193,7 +1193,7 @@ def __init__(
         precision=precision,
     )
 
-  @nnx.jit
+  # @nnx.jit
   def encode(
       self, x: jax.Array, return_dict: bool = True
   ) -> Union[FlaxAutoencoderKLOutput, Tuple[FlaxDiagonalGaussianDistribution]]:

Original file line number	Diff line number	Diff line change
`@@ -1193,7 +1193,7 @@ def __init__(`
`1193`	`1193`	`precision=precision,`
`1194`	`1194`	`)`
`1195`	`1195`
`1196`		`- @nnx.jit`
	`1196`	`+ # @nnx.jit`
`1197`	`1197`	`def encode(`
`1198`	`1198`	`self, x: jax.Array, return_dict: bool = True`
`1199`	`1199`	`) -> Union[FlaxAutoencoderKLOutput, Tuple[FlaxDiagonalGaussianDistribution]]:`