Merge pull request #14559 from Nuullll/ipex-sdpa-fix

[IPEX] Fix SDPA attn_mask dtype
AUTOMATIC1111 · Jan 6, 2024 · b00b429 · b00b429
2 parents 8b6848c + 16b4d2c
commit b00b429
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/modules/xpu_specific.py b/modules/xpu_specific.py
@@ -41,6 +41,8 @@ def torch_xpu_scaled_dot_product_attention(
     # cast to same dtype first
     key = key.to(query.dtype)
     value = value.to(query.dtype)
+    if attn_mask is not None and attn_mask.dtype != torch.bool:
+        attn_mask = attn_mask.to(query.dtype)
 
     N = query.shape[:-2]  # Batch size
     L = query.size(-2)  # Target sequence length