diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py
index 988cd683f..29fabbb33 100644
--- a/fastvideo/v1/pipelines/composed_pipeline_base.py
+++ b/fastvideo/v1/pipelines/composed_pipeline_base.py
@@ -245,7 +245,10 @@ class ComposedPipelineBase(ABC):
                 logger.info("Using module %s already provided", module_name)
                 modules[module_name] = loaded_modules[module_name]
                 continue
-            loading_module_name = module_name.split("_")[-1]
+            if 'transformer' in module_name:
+                loading_module_name = module_name.split("_")[-1]
+            else:
+                loading_module_name = module_name
             component_model_path = os.path.join(self.model_path, loading_module_name)
             # torch.distributed.breakpoint()  
             module = PipelineComponentLoader.load_module(
diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py
index 0904117d4..bc6b3526c 100644
--- a/fastvideo/v1/pipelines/stages/decoding.py
+++ b/fastvideo/v1/pipelines/stages/decoding.py
@@ -103,7 +103,7 @@ class DecodingStage(PipelineStage):
                 #     self.vae.enable_parallel()
                 if not vae_autocast_enabled:
                     latents = latents.to(vae_dtype)
-                torch.distributed.breakpoint()
+                    
                 image = self.vae.decode(latents)
 
         # Normalize image to [0, 1] range
diff --git a/fastvideo/v1/training/distillation_pipeline.py b/fastvideo/v1/training/distillation_pipeline.py
index 9c3003ecc..add1b2d7a 100644
--- a/fastvideo/v1/training/distillation_pipeline.py
+++ b/fastvideo/v1/training/distillation_pipeline.py
@@ -12,7 +12,6 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 import torchvision
-# from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.optimization import get_scheduler
 from einops import rearrange
 from torchdata.stateful_dataloader import StatefulDataLoader
@@ -389,7 +388,7 @@ class DistillationPipeline(TrainingPipeline):
             noise=critic_noise.flatten(0, 1),
             flow_pred=flow_pred
         )
-
+        torch.distributed.breakpoint()
         critic_log_dict = {
             "critictrain_latent": generated_video.detach(),
             "critictrain_noisy_latent": noisy_generated_video.detach(),
@@ -717,7 +716,7 @@ class DistillationPipeline(TrainingPipeline):
                 dmd_latents_name = ['dmdtrain_latents', 'dmdtrain_noisy_latent', 'dmdtrain_pred_real_video', 'dmdtrain_pred_fake_video']
                 for latent_key in dmd_latents_name:
                     latents = generator_log_dict[latent_key]
-                    decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args)
+                    decoded_latent = decode_stage.forward(ForwardBatch(data_type="video", latents=latents), training_args)
                     wandb_loss_dict.update({
                         latent_key: prepare_for_saving(decoded_latent.output),
                     })