sharding before call to vae encoder

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -529,6 +529,11 @@ def prepare_latents_i2v_base(
 
       vae_dtype = getattr(self.vae, "dtype", jnp.float32)
       video_condition = video_condition.astype(vae_dtype)
+      sharding_spec = P(self.config.mesh_axes[0], None, None, None, None)
+      video_condition = jax.lax.with_sharding_constraint(
+          video_condition, 
+          sharding_spec
+      )
 
       with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
           encoded_output = self.vae.encode(video_condition, self.vae_cache)[0].mode()