diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py
index 6e50d3bc..f6b544a8 100644
--- a/extensions_built_in/diffusion_models/ltx2/ltx2.py
+++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py
@@ -838,7 +838,7 @@ class LTX2Model(BaseModel):
                     batch_size,
                     num_channels_latents=num_channels_latents_audio,
                     num_mel_bins=num_mel_bins,
-                    num_frames=batch.tensor.shape[1],
+                    num_frames=batch.dataset_config.num_frames,
                     frame_rate=frame_rate,
                     sampling_rate=self.pipeline.audio_sampling_rate,
                     hop_length=self.pipeline.audio_hop_length,