fix

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_transformers.py

@@ -7,6 +7,7 @@
 from onnx_diagnostic.ext_test_case import (
     ExtTestCase,
     requires_transformers,
+    requires_torch,
     ignore_warnings,
     has_onnxscript,
 )
@@ -465,6 +466,7 @@ def forward(
             )
 
     @requires_transformers("4.99")
+    @requires_torch("2.9.99")
     @unittest.skipIf(not patch_qwen2_5, "Qwen25 not part of this transformers")
     def test_qwen2_5_vl_vision_attention_iteration(self):
         from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -2435,8 +2435,8 @@ def _iteration(start_end, query_states, key_states, value_states):
                         cu_seqlens.dtype
                     )
                     dot = dot.sum(dim=0)
-                    mask = dot.unsqueeze(1) @ dot.unsqueeze(0)
-                    bool_mask = mask == dot**2
+                    mask = dot.unsqueeze(1) - dot.unsqueeze(0)
+                    bool_mask = mask == 0
                     bool_mask = bool_mask.unsqueeze(0).unsqueeze(0)
 
                     torch._check(bool_mask.shape[2] == key_states.shape[2])