changes to transformer

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -665,14 +665,16 @@ def __init__(
     self.caption_projection = NNXPixArtAlphaTextProjection(
         rngs=rngs,
         in_features=self.caption_channels,
-        hidden_size=inner_dim,
+        hidden_size=self.caption_channels,
+        out_features=inner_dim,
         dtype=self.dtype,
         weights_dtype=self.weights_dtype,
     )
     self.audio_caption_projection = NNXPixArtAlphaTextProjection(
         rngs=rngs,
         in_features=self.caption_channels,
-        hidden_size=audio_inner_dim,
+        hidden_size=self.caption_channels,
+        out_features=audio_inner_dim,
         dtype=self.dtype,
         weights_dtype=self.weights_dtype,
     )