Merge pull request #33 from GavChap/sdp-attn-fix

Update PixArt_blocks.py
This commit is contained in:
City
2024-05-12 22:07:52 +02:00
committed by GitHub
+1
View File
@@ -189,6 +189,7 @@ class AttentionKVCompress(Attention_):
q, k, v = map(lambda t: t.transpose(1, 2),(q, k, v),) q, k, v = map(lambda t: t.transpose(1, 2),(q, k, v),)
x = torch.nn.functional.scaled_dot_product_attention( x = torch.nn.functional.scaled_dot_product_attention(
q, k, v, q, k, v,
dropout_p=self.attn_drop.p,
attn_mask=attn_bias attn_mask=attn_bias
).transpose(1, 2).contiguous() ).transpose(1, 2).contiguous()
x = x.view(B, N, C) x = x.view(B, N, C)