Merge branch 'attn_mask_fix' of https://github.com/ebsmothers/text into attn_mask_fix

Author: ebsmothers

test/models/test_models.py

@@ -15,8 +15,8 @@ def test_self_attn_mask(self):
         query[0, ...] = 0
         key_padding_mask = torch.zeros((batch_size, source_len))
         float_attn_mask = torch.zeros((source_len, source_len))
-        bool_attn_mask = float_attn_mask.to(dtype=bool)
         float_attn_mask[0][1] = -1e8
+        bool_attn_mask = float_attn_mask.to(dtype=bool)
         with torch.no_grad():
             mha.input_projection.weight.fill_(1. / embed_dim)
             mha.input_projection.bias.fill_(0.)

torchtext/models/roberta/modules.py

@@ -286,6 +286,7 @@ def __init__(
     def forward(self, tokens: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> Union[torch.Tensor, List[torch.Tensor]]:
         if attn_mask is not None:
             torch._assert(attn_mask.is_floating_point() or attn_mask.dtype == torch.bool, f"Only float or bool types are supported for attn_mask not {attn_mask.dtype}")
+
         padding_mask = tokens.eq(self.padding_idx)
 
         token_embeddings = self.token_embedding(tokens)