disable shardy for generate_wan

jfacevedo-google · jfacevedo-google · commit 500d1c1d5a8c · 2025-07-11T15:52:39.000Z
diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py
@@ -20,7 +20,7 @@
 from absl import app
 from maxdiffusion.utils import export_to_video
 
-jax.config.update("jax_use_shardy_partitioner", True)
+jax.config.update("jax_use_shardy_partitioner", False)
 
 
 def run(config, pipeline=None, filename_prefix=""):
