diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
index d1d9937e3..a7e0fca7d 100644
--- a/src/maxdiffusion/models/attention_flax.py
+++ b/src/maxdiffusion/models/attention_flax.py
@@ -3324,8 +3324,6 @@ def __init__(
     if attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
 
-      jax.config.update("jax_use_shardy_partitioner", False)
-
       dpa_layer = DotProductAttention(
           head_dim=dim_head,
           num_attention_heads=heads,
@@ -3437,8 +3435,6 @@ def setup(self):
     if self.attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
 
-      jax.config.update("jax_use_shardy_partitioner", False)
-
       dpa_layer = DotProductAttention(
           head_dim=self.dim_head,
           num_attention_heads=self.heads,