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

Update PixArt_blocks.py
city96 · May 12, 2024 · 1e5dc8e · 1e5dc8e
2 parents c2e2f6e + 3824b0d
commit 1e5dc8e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/PixArt/models/PixArt_blocks.py b/PixArt/models/PixArt_blocks.py
@@ -189,6 +189,7 @@ def forward(self, x, mask=None, HW=None, block_id=None):
             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 @@ def forward(self, caption, train, force_drop_ids=None):
         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
+        return y_embed, caption