sharding before call to vae encoder

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -529,13 +529,12 @@ 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):
+          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