[Bugfix] fix _normalize_dit_input (#681)

2025-08-01 05:10:27 -04:00
parent 2c53d3eecf
commit fcba463553
1 changed files with 3 additions and 3 deletions
@@ -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,