removing nnx.jit decorator

prishajain1 · prishajain1 · commit ce8ab0d6614c · 2026-02-26T22:08:07.000+05:30
diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
@@ -774,7 +774,6 @@ def __init__(
         precision=precision,
     )
 
-  @nnx.jit(static_argnames="feat_idx")
   def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
     if feat_cache is not None:
       idx = feat_idx
@@ -923,7 +922,6 @@ def __init__(
         precision=precision,
     )
 
-  @nnx.jit(static_argnames="feat_idx")
   def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
     if feat_cache is not None:
       idx = feat_idx

Original file line number	Diff line number	Diff line change
`@@ -774,7 +774,6 @@ def __init__(`
`774`	`774`	`precision=precision,`
`775`	`775`	`)`
`776`	`776`
`777`		`- @nnx.jit(static_argnames="feat_idx")`
`778`	`777`	`def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):`
`779`	`778`	`if feat_cache is not None:`
`780`	`779`	`idx = feat_idx`
`@@ -923,7 +922,6 @@ def __init__(`
`923`	`922`	`precision=precision,`
`924`	`923`	`)`
`925`	`924`
`926`		`- @nnx.jit(static_argnames="feat_idx")`
`927`	`925`	`def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):`
`928`	`926`	`if feat_cache is not None:`
`929`	`927`	`idx = feat_idx`