fixed num_attention_heads in test

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -268,16 +268,16 @@ def test_transformer_3d_model_dot_product_attention(self):
             out_channels=self.out_channels,
             patch_size=self.patch_size,
             patch_size_t=self.patch_size_t,
-            num_attention_heads=self.num_attention_heads,
-            attention_head_dim=self.attention_head_dim,
-            cross_attention_dim=self.cross_attention_dim,
+            num_attention_heads=8,
+            attention_head_dim=128,
+            cross_attention_dim=32,
             audio_in_channels=self.audio_in_channels,
-            audio_out_channels=self.audio_out_channels,
+            audio_out_channels=self.audio_in_channels,
             audio_patch_size=self.audio_patch_size,
             audio_patch_size_t=self.audio_patch_size_t,
-            audio_num_attention_heads=self.audio_num_attention_heads,
-            audio_attention_head_dim=self.audio_attention_head_dim,
-            audio_cross_attention_dim=self.audio_cross_attention_dim,
+            audio_num_attention_heads=8,
+            audio_attention_head_dim=128,
+            audio_cross_attention_dim=32,
             num_layers=1, # Reduced layers for speed
             config=self.config,
             scan_layers=False,