update util to work for additive attention mask

facebookresearch · Nov 8, 2021 · 8c6d250 · 8c6d250
1 parent 8f3c01e
commit 8c6d250
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/xformers/components/attention/utils.py b/xformers/components/attention/utils.py
@@ -49,7 +49,10 @@ def maybe_merge_masks(
         if att_mask is None:
             att_mask = key_padding_mask
         # Assumption is that False means to mask.
-        att_mask = att_mask.logical_and(key_padding_mask)
+        elif att_mask.dtype == torch.bool:
+            att_mask = att_mask.logical_and(key_padding_mask)
+        else:
+            att_mask = att_mask.masked_fill(~key_padding_mask, float("-inf"))
 
     return att_mask