Masking: Mask everything after a specified token (eos)
My tgt tensor is in shape of [12, 32, 1] which is sequence_length, batch_size, token_idx.

What is the best way to create a mask which has ones for entries with &lt;eos&gt; and before in sequence, and zeros afterwards?

Currently I'm calculating my mask like this, which simply puts zeros where &lt;blank&gt; is, ones otherwise.

mask = torch.zeros_like(tgt).masked_scatter_((tgt != tgt_padding), torch.ones_like(tgt))

But the problem is, that my tgt can contain &lt;blank&gt; as well (before &lt;eos&gt;), in which cases I don't want to mask it out.

My temporary solution:

mask = torch.ones_like(tgt)
for eos_token in (tgt == tgt_eos).nonzero():
    mask[eos_token[0]+1:,eos_token[1]] = 0


I guess you are trying to create a mask for the PAD tokens. There are several ways. One of them is as follows.

# tensor is of shape [seq_len, batch_size, 1]
tensor = tensor.mul(tensor.ne(PAD).float())


Here, PAD stands for the index of the PAD_TOKEN. tensor.ne(PAD) will create a byte tensor where at PAD_TOKEN positions, 0 will be assigned and 1 elsewhere.



If you have examples like, "&lt;s&gt; I think &lt;pad&gt; so &lt;/s&gt; &lt;pad&gt; &lt;pad&gt;". Then, I would suggest using different PAD tokens, for before and after &lt;/s&gt;. 

OR, if you have the length information for each sentence (in the above example, the sentence length is 6), then you can create the mask using the following function.

def sequence_mask(lengths, max_len=None):
    """
    Creates a boolean mask from sequence lengths.
    :param lengths: 1d tensor [batch_size]
    :param max_len: int
    """
    batch_size = lengths.numel()
    max_len = max_len or lengths.max()
    return (torch.arange(0, max_len, device=lengths.device)  # (0 for pad positions)
            .type_as(lengths)
            .repeat(batch_size, 1)
            .lt(lengths.unsqueeze(1)))

