diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py
index 18ebb05..450a046 100644
--- a/src/h3_blackwell_runtime/vae_decoder.py
+++ b/src/h3_blackwell_runtime/vae_decoder.py
@@ -316,7 +316,7 @@ class MiniMaxH3VideoVAE(nn.Module):
             pad_tokens += self.tokens_chunk_size
             chunks += 1
         if pad_tokens:
-            z = torch.cat((z, z[:, :, -1:].expand(-1, -1, pad_tokens, -1, -1)), dim=2)
+            z = torch.cat((z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)), dim=2)
         output_frames = self._decode_temporal_frame_plan(z.shape[2], chunks, pad_tokens)
         output = None
         overlap = None