only step jit

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

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