Changed none -> data

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -151,7 +151,7 @@ def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) ->
       # Shape is (Batch, Time, Height, Width, Channels)
       # We only shard if the dimension is divisible by the mesh size to avoid XLA errors
       if x_padded.shape[2] % self.mesh.shape["context"] == 0:
-        sharding = NamedSharding(self.mesh, P(None, None, "context", None, None))
+        sharding = NamedSharding(self.mesh, P("data", None, "context", None, None))
         x_padded = jax.lax.with_sharding_constraint(x_padded, sharding)
 
     out = self.conv(x_padded)