From 272fcf2b38121331a20edb06be89a0fa2ba2f06b Mon Sep 17 00:00:00 2001 From: Markus Schnoes Date: Fri, 18 Oct 2024 13:28:02 +0200 Subject: [PATCH] [PyTorch] Use or instead of and to combine swa mask with existing mask Signed-off-by: Markus Schnoes --- transformer_engine/pytorch/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py index 62ffec2cd6..5ca525ce97 100644 --- a/transformer_engine/pytorch/attention.py +++ b/transformer_engine/pytorch/attention.py @@ -1075,7 +1075,7 @@ def get_swa_mask( attn_mask_type = "arbitrary" mask = mask_lower.logical_not() if attention_mask is not None: - mask = torch.logical_and(attention_mask, mask) + mask = torch.logical_or(attention_mask, mask) return attn_mask_type, mask