Update PixArt_blocks.py

Add dropout_p to SDP attention
2024-05-11 17:00:48 +01:00
parent c2e2f6e810
commit 3824b0d7b9
1 changed files with 2 additions and 1 deletions
@@ -189,6 +189,7 @@ class AttentionKVCompress(Attention_):
            q, k, v = map(lambda t: t.transpose(1, 2),(q, k, v),)
            x = torch.nn.functional.scaled_dot_product_attention(
                q, k, v,
+                dropout_p=self.attn_drop.p,
                attn_mask=attn_bias
            ).transpose(1, 2).contiguous()
        x = x.view(B, N, C)