dim changed in test

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -272,14 +272,14 @@ def test_scan_remat_parity(self):
             patch_size=self.patch_size,
             patch_size_t=self.patch_size_t,
             num_attention_heads=8,
-            attention_head_dim=4,
+            attention_head_dim=128,
             num_layers=2, # Need >1 layer to test scan effectively
             caption_channels=32,
             cross_attention_dim=32,
             audio_in_channels=self.audio_in_channels,
             audio_out_channels= self.audio_in_channels,
             audio_num_attention_heads=8,
-            audio_attention_head_dim=4,
+            audio_attention_head_dim=128,
             audio_cross_attention_dim=32
         )