Refactor

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -136,7 +136,7 @@ def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) ->
 
     if self.mesh is not None:
       # (B, D, H, W, C)
-      x_padded = with_sharding_constraint(x_padded, PartitionSpec(None, None, 'fsdp', None, None))
+      x_padded = with_sharding_constraint(x_padded, PartitionSpec(None, None, None, 'fsdp', None))
 
     out = self.conv(x_padded)
     return out