bed0a0ae26451c9897cf1ee0f7302e42eba9b42c,transformer/Models.py,,get_attn_subsequent_mask,#,31

Before Change

    """ Get an attention mask to avoid using the subsequent info."""

    assert seq.dim() == 2
    attn_shape = (seq.size(0), seq.size(1), seq.size(1))
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype("uint8")
    subsequent_mask = torch.from_numpy(subsequent_mask)
    if seq.is_cuda:
        subsequent_mask = subsequent_mask.cuda()
    return subsequent_mask

class Encoder(nn.Module):
    """ A encoder model with self attention mechanism. """

After Change

    sz_b, len_s = seq.size()

    subsequent_mask = torch.ones((len_s, len_s), device=seq.device, dtype=torch.uint8)
    subsequent_mask = torch.triu(subsequent_mask, diagonal=1)
    subsequent_mask = subsequent_mask.unsqueeze(0).expand(sz_b, len_s, len_s)
    return subsequent_mask

Instances