spell

Author: xadupre

_doc/recipes/plot_export_dim1.py

@@ -75,8 +75,8 @@ def forward(self, x, y, z):
 
 
 # %%
-# Final try with pathes...
-# ++++++++++++++++++++++++
+# Final try with patches...
+# +++++++++++++++++++++++++
 
 print("-- export shape:", string_type((x, y, z), with_shape=True))
 print("-- dynamic shapes:", string_type((ds, ds, ds)))

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1351,10 +1351,6 @@ def patched_sdpa_attention_forward(
         "`sdpa` attention does not support `output_attentions=True`."
         " Please set your attention to `eager` if you want any of these features."
     )
-    torch._check(
-        attention_mask is None or attention_mask.shape[3] == key.shape[2],
-        "Attention mask shape incompatible with key shape.",
-    )
     torch._check(
         query.shape[0] == key.shape[0] or query.shape[0] == 1,
         lambda: (
@@ -1385,6 +1381,11 @@ def patched_sdpa_attention_forward(
     if attention_mask is not None and attention_mask.ndim == 4:
         attention_mask = attention_mask[:, :, :, : key.shape[-2]]
 
+    torch._check(
+        attention_mask is None or attention_mask.shape[3] == key.shape[2],
+        lambda: "Attention mask shape incompatible with key shape.",
+    )
+
     if patch_sdpa_is_causal:
         # transformers>=4.55
         is_causal = is_causal if is_causal is not None else getattr(module, "is_causal", True)