source sync

Google-ML-Automation · michelle-yooh · commit cd5c94da428c · 2026-02-06T19:00:32.000Z
PiperOrigin-RevId: 866530641
diff --git a/src/maxdiffusion/__init__.py b/src/maxdiffusion/__init__.py
@@ -67,7 +67,7 @@
 if is_flax_available():
   from flax import config as flax_config
 
-  flax_config.update('flax_always_shard_variable', False)
+  flax_config.update("flax_always_shard_variable", False)
 
 try:
   if not is_onnx_available():
