diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py
index a4c7dd7f5..ae9e14a88 100644
--- a/fastvideo/v1/training/training_pipeline.py
+++ b/fastvideo/v1/training/training_pipeline.py
@@ -233,6 +233,11 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         training_batch.sigmas = sigmas
         training_batch.noise = noise
 
+        # shard noisy_model_input to match
+        training_batch.noisy_model_input = shard_latents_across_sp(
+            training_batch.noisy_model_input,
+            num_latent_t=self.training_args.num_latent_t)
+
         return training_batch
 
     def _build_attention_metadata(
@@ -367,10 +372,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             training_batch.latents = shard_latents_across_sp(
                 training_batch.latents,
                 num_latent_t=self.training_args.num_latent_t)
-            # shard noisy_model_input to match
-            training_batch.noisy_model_input = shard_latents_across_sp(
-                training_batch.noisy_model_input,
-                num_latent_t=self.training_args.num_latent_t)
             # shard noise to match latents
             training_batch.noise = shard_latents_across_sp(
                 training_batch.noise,