Removed unused imports

prishajain1 · prishajain1 · commit 89c5ebdb0e93 · 2026-01-26T20:31:26.000+05:30
diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
@@ -39,11 +39,6 @@
 except LookupError:
   pass
 
-
-from jax.sharding import PartitionSpec
-from jax.lax import with_sharding_constraint
-
-
 def _update_cache(cache, idx, value):
   if cache is None:
     return None
