fix patch

xadupre · xadupre · commit ae0b77c39a34 · 2025-10-23T13:03:11.000+02:00
diff --git a/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py b/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py
@@ -1313,7 +1313,7 @@ def patched_sdpa_attention_forward(
     is_causal = attention_mask is None and is_causal
 
     torch._check(
-        attention_mask.shape[3] == key.shape[2],
+        attention_mask is None or attention_mask.shape[3] == key.shape[2],
         "Attention mask shape incompatible with key shape.",
     )
     attn_output = torch.nn.functional.scaled_dot_product_attention(

Original file line number	Diff line number	Diff line change
`@@ -1313,7 +1313,7 @@ def patched_sdpa_attention_forward(`
`1313`	`1313`	`is_causal = attention_mask is None and is_causal`
`1314`	`1314`
`1315`	`1315`	`torch._check(`
`1316`		`- attention_mask.shape[3] == key.shape[2],`
	`1316`	`+ attention_mask is None or attention_mask.shape[3] == key.shape[2],`
`1317`	`1317`	`"Attention mask shape incompatible with key shape.",`
`1318`	`1318`	`)`
`1319`	`1319`	`attn_output = torch.nn.functional.scaled_dot_product_attention(`