fixing hunyuan attn mask test case

dabeschte · dabeschte · commit 22c86731d717 · 2025-02-11T10:11:59.000Z
diff --git a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py
@@ -204,4 +204,7 @@ def prepare_causal_attention_mask_orig(
         new_mask = prepare_causal_attention_mask(
             num_frames=31, height_width=111, dtype=torch.float32, device=torch_device
         )
-        self.assertEqual(original_mask, new_mask)
+        self.assertTrue(
+            torch.allclose(original_mask, new_mask),
+            "Causal attention mask should be the same",
+        )

Original file line number	Diff line number	Diff line change
`@@ -204,4 +204,7 @@ def prepare_causal_attention_mask_orig(`
`204`	`204`	`new_mask = prepare_causal_attention_mask(`
`205`	`205`	`num_frames=31, height_width=111, dtype=torch.float32, device=torch_device`
`206`	`206`	`)`
`207`		`- self.assertEqual(original_mask, new_mask)`
	`207`	`+ self.assertTrue(`
	`208`	`+ torch.allclose(original_mask, new_mask),`
	`209`	`+ "Causal attention mask should be the same",`
	`210`	`+ )`