From 3824b0d7b9cbee04ac32c2bd6f199da3fc78f6b7 Mon Sep 17 00:00:00 2001 From: GavChap Date: Sat, 11 May 2024 17:00:48 +0100 Subject: [PATCH] Update PixArt_blocks.py Add dropout_p to SDP attention --- PixArt/models/PixArt_blocks.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/PixArt/models/PixArt_blocks.py b/PixArt/models/PixArt_blocks.py index a7b604b..2fbc176 100644 --- a/PixArt/models/PixArt_blocks.py +++ b/PixArt/models/PixArt_blocks.py @@ -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) @@ -467,4 +468,4 @@ class CaptionEmbedderDoubleBr(nn.Module): if (train and use_dropout) or (force_drop_ids is not None): global_caption, caption = self.token_drop(global_caption, caption, force_drop_ids) y_embed = self.proj(global_caption) - return y_embed, caption \ No newline at end of file + return y_embed, caption