mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-02 22:27:18 +08:00
Update attention.py
Fix https://github.com/comfyanonymous/ComfyUI/issues/8634
This commit is contained in:
parent
2a0b138feb
commit
9a33651bb5
@ -446,6 +446,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
# add a heads dimension if there isn't already one
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
mask = mask.contiguous()
|
||||
|
||||
if SDP_BATCH_LIMIT >= b:
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user