From fcba463553a6dae679d775d3a8a28a8f9d3d1263 Mon Sep 17 00:00:00 2001
From: MartinPernus <36865527+MartinPernus@users.noreply.github.com>
Date: Fri, 1 Aug 2025 11:10:27 +0200
Subject: [PATCH] [Bugfix] fix _normalize_dit_input (#681)

---
 fastvideo/training/training_pipeline.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py
index 0e23783ed..12b457aac 100644
--- a/fastvideo/training/training_pipeline.py
+++ b/fastvideo/training/training_pipeline.py
@@ -209,9 +209,9 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _normalize_dit_input(self,
                              training_batch: TrainingBatch) -> TrainingBatch:
         # TODO(will): support other models
-        training_batch.latents = normalize_dit_input(
-            'wan', training_batch.latents,
-            self.validation_pipeline.get_module("vae"))
+        training_batch.latents = normalize_dit_input('wan',
+                                                     training_batch.latents,
+                                                     self.get_module("vae"))
         return training_batch
 
     def _prepare_dit_inputs(self,