diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py
index c67b20a13..493baeb81 100644
--- a/fastvideo/models/utils.py
+++ b/fastvideo/models/utils.py
@@ -5,9 +5,6 @@ from typing import Any
 
 import torch
 
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
 
 # TODO(PY): move it elsewhere
 def auto_attributes(init_func):
diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py
index 735bfc7ad..bf971231f 100644
--- a/fastvideo/pipelines/composed_pipeline_base.py
+++ b/fastvideo/pipelines/composed_pipeline_base.py
@@ -84,7 +84,6 @@ class ComposedPipelineBase(ABC):
 
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
-        # fastvideo_args.dit_cpu_offload = False
         with self.profiler_controller.region("profiler_region_model_loading"):
             self.modules = self.load_modules(fastvideo_args, loaded_modules)
 
diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py
index 51bee7836..4d40dec4d 100644
--- a/fastvideo/pipelines/stages/denoising.py
+++ b/fastvideo/pipelines/stages/denoising.py
@@ -241,9 +241,12 @@ class DenoisingStage(PipelineStage):
             temporal_scale = fastvideo_args.pipeline_config.vae_config.temporal_compression_ratio
             spatial_scale = fastvideo_args.pipeline_config.vae_config.spatial_compression_ratio
             patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
-            seq_len = ((F - 1) // temporal_scale + 1) * (
-                batch.height // spatial_scale) * (
-                    batch.width // spatial_scale) // (patch_size * patch_size)
+            if isinstance(patch_size, int):
+                patch_size = (patch_size, patch_size, patch_size)
+            seq_len = ((F - 1) // temporal_scale +
+                       1) * (batch.height // spatial_scale) * (
+                           batch.width // spatial_scale) // (patch_size[1] *
+                                                             patch_size[2])
 
         # Initialize lists for ODE trajectory
         trajectory_timesteps: list[torch.Tensor] = []
@@ -414,7 +417,6 @@ class DenoisingStage(PipelineStage):
                             **action_kwargs,
                         )
 
-                    assert batch.do_classifier_free_guidance, "do_classifier_free_guidance is not supported"
                     if batch.do_classifier_free_guidance:
                         batch.is_cfg_negative = True
                         with set_forward_context(
diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py
index adef13acd..32ef93cf2 100644
--- a/fastvideo/pipelines/stages/input_validation.py
+++ b/fastvideo/pipelines/stages/input_validation.py
@@ -118,10 +118,11 @@ class InputValidationStage(PipelineStage):
                 oh, ow = batch.height, batch.width
                 img = img.resize((ow, oh), Image.LANCZOS)
             else:
-                # Standard Wan logic
                 patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+                if isinstance(patch_size, int):
+                    patch_size = (patch_size, patch_size, patch_size)
                 vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
-                dh, dw = patch_size * vae_stride, patch_size * vae_stride
+                dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
                 max_area = 480 * 848
                 ow, oh = best_output_size(iw, ih, dw, dh, max_area)
 
diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py
index 289573ebc..491c39a6d 100644
--- a/fastvideo/training/self_forcing_distillation_pipeline.py
+++ b/fastvideo/training/self_forcing_distillation_pipeline.py
@@ -67,7 +67,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         self.noise_scheduler = SelfForcingFlowMatchScheduler(
             num_inference_steps=1000,
-            shift=5.0,
+            shift=self.timestep_shift,
             sigma_min=0.0,
             extra_one_step=True,
             training=True)