Commit 16b4d2cf authored by Nuullll's avatar Nuullll

[IPEX] Fix SDPA attn_mask dtype

parent 8b6848c6
...@@ -41,6 +41,8 @@ def torch_xpu_scaled_dot_product_attention( ...@@ -41,6 +41,8 @@ def torch_xpu_scaled_dot_product_attention(
# cast to same dtype first # cast to same dtype first
key = key.to(query.dtype) key = key.to(query.dtype)
value = value.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 N = query.shape[:-2] # Batch size
L = query.size(-2) # Target sequence length L = query.size(-2) # Target sequence length
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment