Change in wan_vae_test.py

Author: prishajain1

src/maxdiffusion/tests/wan_vae_test.py

@@ -522,11 +522,11 @@ def vae_encode(video, wan_vae, vae_cache, key):
     params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), params)
     wan_vae = nnx.merge(graphdef, params)
 
-    p_vae_encode = jax.jit(functools.partial(vae_encode, wan_vae=wan_vae, vae_cache=vae_cache, key=key))
+    p_vae_encode = functools.partial(vae_encode, wan_vae=wan_vae, vae_cache=vae_cache, key=key)
     original_video_shape = original_video.shape
     latent = p_vae_encode(original_video)
 
-    jitted_decode = jax.jit(functools.partial(wan_vae.decode, feat_cache=vae_cache, return_dict=False))
+    jitted_decode = functools.partial(wan_vae.decode, feat_cache=vae_cache, return_dict=False)
     video = jitted_decode(latent)[0]
     video = jnp.transpose(video, (0, 4, 1, 2, 3))
     assert video.shape == original_video_shape