diff --git a/fastvideo/v1/pipelines/wan/wan_pipeline.py b/fastvideo/v1/pipelines/wan/wan_pipeline.py
index 999f7fe36..6ef0444c8 100644
--- a/fastvideo/v1/pipelines/wan/wan_pipeline.py
+++ b/fastvideo/v1/pipelines/wan/wan_pipeline.py
@@ -30,6 +30,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # We use UniPCMScheduler from Wan2.1 official repo, not the one in diffusers. 
         self.modules["scheduler"] = FlowUniPCMultistepScheduler(
             shift=fastvideo_args.flow_shift)
 
diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py
index 48340e87b..4c3afdf16 100644
--- a/fastvideo/v1/training/training_pipeline.py
+++ b/fastvideo/v1/training/training_pipeline.py
@@ -194,6 +194,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
 
         # Add the transformer to the validation pipeline
         self.validation_pipeline.add_module("transformer", transformer)
+        # TODO(Peiyuan): those logic should be inside add_module
         self.validation_pipeline.latent_preparation_stage.transformer = transformer  # type: ignore[attr-defined]
         self.validation_pipeline.denoising_stage.transformer = transformer  # type: ignore[attr-defined]
 
@@ -237,7 +238,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             )
 
             # Run validation inference
-            with torch.inference_mode(), torch.autocast("cuda",
+            with torch.no_grad(), torch.autocast("cuda",
                                                         dtype=torch.bfloat16):
                 output_batch = self.validation_pipeline.forward(
                     batch, training_args)
diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py
index 02d1d198d..e92c518f0 100644
--- a/fastvideo/v1/training/wan_training_pipeline.py
+++ b/fastvideo/v1/training/wan_training_pipeline.py
@@ -271,7 +271,8 @@ class WanTrainingPipeline(TrainingPipeline):
         gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
         logger.info("GPU memory usage before train_one_step: %s MB",
                     gpu_memory_usage)
-
+        
+        self._log_validation(self.transformer, self.training_args, 1)
         for step in range(self.init_steps + 1,
                           self.training_args.max_train_steps + 1):
             start_time = time.perf_counter()