fix

prishajain1 · prishajain1 · commit 2ed7122b8151 · 2026-02-09T21:23:45.000+05:30
diff --git a/src/maxdiffusion/tests/ltx_2_transformer_test.py b/src/maxdiffusion/tests/ltx_2_transformer_test.py
@@ -313,8 +313,8 @@ def test_transformer_3d_model_dot_product_attention(self):
                 audio_encoder_attention_mask=audio_encoder_attention_mask
             )
             
-        self.assertEqual(output.sample.shape, hidden_states.shape)
-        self.assertEqual(output.audio_sample.shape, audio_hidden_states.shape)
+        self.assertEqual(output["sample"].shape, hidden_states.shape)
+        self.assertEqual(output["audio_sample"].shape, audio_hidden_states.shape)
 
     def test_scan_remat_parity(self):
         """

Original file line number	Diff line number	Diff line change
`@@ -313,8 +313,8 @@ def test_transformer_3d_model_dot_product_attention(self):`
`313`	`313`	`audio_encoder_attention_mask=audio_encoder_attention_mask`
`314`	`314`	`)`
`315`	`315`
`316`		`- self.assertEqual(output.sample.shape, hidden_states.shape)`
`317`		`- self.assertEqual(output.audio_sample.shape, audio_hidden_states.shape)`
	`316`	`+ self.assertEqual(output["sample"].shape, hidden_states.shape)`
	`317`	`+ self.assertEqual(output["audio_sample"].shape, audio_hidden_states.shape)`
`318`	`318`
`319`	`319`	`def test_scan_remat_parity(self):`
`320`	`320`	`"""`