diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py
index 0c18f84d1..9421fddb7 100644
--- a/fastvideo/pipelines/stages/causal_denoising.py
+++ b/fastvideo/pipelines/stages/causal_denoising.py
@@ -349,7 +349,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                             0, 2, 1, 3, 4)
                     else:
                         current_latents = pred_video_btchw.permute(
-                            0, 2, 1, 3, 4)                        
+                            0, 2, 1, 3, 4)
 
                     if progress_bar is not None:
                         progress_bar.update()
@@ -374,6 +374,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                     t_expanded_context = t_context.unsqueeze(1)
 
                     if boundary_timestep is not None:
+                        assert self.transformer_2 is not None, "transformer_2 is not provided, but boundary_timestep is not None"
                         self.transformer_2(
                             context_bcthw,
                             prompt_embeds,
diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py
index 100199ec8..c51642cb4 100644
--- a/fastvideo/pipelines/stages/input_validation.py
+++ b/fastvideo/pipelines/stages/input_validation.py
@@ -5,7 +5,6 @@ Input validation stage for diffusion pipelines.
 
 import torch
 import torchvision.transforms.functional as TF
-from PIL import Image
 
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -14,7 +13,6 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import (StageValidators,
                                                    VerificationResult)
-from fastvideo.utils import best_output_size
 
 logger = init_logger(__name__)