vocoder bwe changes

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_bwe_ltx2.py

@@ -489,7 +489,7 @@ def __call__(self, y: Array) -> Tuple[Array, Array]:
     w_transposed = jnp.transpose(self.forward_basis.value, (2, 1, 0)) # (O, I, K)
 
     spec = jax.lax.conv_general_dilated(
-        lhs=y_transposed,
+        lhs=y_transposed.astype(w_transposed.dtype),
         rhs=w_transposed,
         window_strides=(self.hop_length,),
         padding="VALID",