test file dims changed

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -184,12 +184,12 @@ def test_transformer_3d_model_instantiation_and_forward(self):
                 attention_head_dim=128,
                 num_layers=1, # 1 layer for speed
                 caption_channels=32, # small for test
-                cross_attention_dim=32,
+                cross_attention_dim=1024,
                 audio_in_channels=self.audio_in_channels,
                 audio_out_channels= self.audio_in_channels,
                 audio_num_attention_heads=8,
                 audio_attention_head_dim=128,
-                audio_cross_attention_dim=32,
+                audio_cross_attention_dim=1024,
                 mesh=self.mesh,
             )
 
@@ -270,15 +270,15 @@ def test_transformer_3d_model_dot_product_attention(self):
             patch_size_t=self.patch_size_t,
             num_attention_heads=8,
             attention_head_dim=128,
-            cross_attention_dim=32,
+            cross_attention_dim=1024,
             caption_channels=32,
             audio_in_channels=self.audio_in_channels,
             audio_out_channels=self.audio_in_channels,
             audio_patch_size=1,
             audio_patch_size_t=1,
             audio_num_attention_heads=8,
             audio_attention_head_dim=128,
-            audio_cross_attention_dim=32,
+            audio_cross_attention_dim=1024,
             num_layers=1, # Reduced layers for speed
             scan_layers=False,
             mesh=self.mesh,