Adding check for batch size divisibility before sharding video condition tensor

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -544,8 +544,11 @@ def prepare_latents_i2v_base(
     vae_dtype = getattr(self.vae, "dtype", jnp.float32)
     video_condition = video_condition.astype(vae_dtype)
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-      sharding_spec = P(self.config.mesh_axes[0], None, None, None, None)
-      video_condition = jax.lax.with_sharding_constraint(video_condition, sharding_spec)
+      data_axis_name = self.config.mesh_axes[0]
+      data_mesh_size = self.mesh.shape[self.config.mesh_axes[0]]
+      if video_condition.shape[0] % data_mesh_size == 0:
+        sharding_spec = P(self.config.mesh_axes[0], None, None, None, None)
+        video_condition = jax.lax.with_sharding_constraint(video_condition, sharding_spec)
       encoded_output = self.vae.encode(video_condition, self.vae_cache)[0].mode()
 
     # Normalize latents