Fix reshape size mismatch in BWE vocoder

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -637,7 +637,7 @@ def __call__(self, mel_spec: Array) -> Array:
     if remainder != 0:
       x = jnp.pad(x, ((0, 0), (0, self.hop_length - remainder), (0, 0)))
 
-    x_flattened = x.transpose(0, 2, 1).reshape(-1, num_samples, 1)
+    x_flattened = x.transpose(0, 2, 1).reshape(-1, x.shape[1], 1)
     log_mel, _, _, _ = self.mel_stft(x_flattened)
     log_mel = log_mel.reshape(batch_size, num_channels, -1, log_mel.shape[-1])