update util to work for additive attention mask (#86)

Co-authored-by: Diana Liskovich <dianaml@devfair0471.h2.fair>
facebookresearch · Nov 9, 2021 · cf46d43 · cf46d43
1 parent 03d58e1
commit cf46d43
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