diff --git a/fastvideo/training/wan_i2v_training_pipeline.py b/fastvideo/training/wan_i2v_training_pipeline.py
index 18810b0c1..9f07191ba 100644
--- a/fastvideo/training/wan_i2v_training_pipeline.py
+++ b/fastvideo/training/wan_i2v_training_pipeline.py
@@ -27,7 +27,7 @@ class WanI2VTrainingPipeline(TrainingPipeline):
     """
     A training pipeline for Wan.
     """
-    _required_config_modules = ["scheduler", "transformer"]
+    _required_config_modules = ["scheduler", "transformer", "vae"]
 
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         self.modules["scheduler"] = FlowUniPCMultistepScheduler(