Adding tgt_mask or tgt_key_padding_mask gives nan output

misharuko: <p>I have a basic transformer model:</p>
<pre><code class="lang-auto">class model(nn.Module):
    def __init__(self):
        super().__init__()
        self.seqLength = 5
        self.d_model = 3
        self.nhead = 1
        self.num_encoder_layers = 6
        self.num_decoder_layers = 6
        self.dim_feedforward = 1024
        self.dropout = 0.1
        self.activation = F.relu
        self.batch_first = True

        self.transformer = nn.Transformer(
                d_model = self.d_model,
                nhead = self.nhead,
                num_encoder_layers = self.num_encoder_layers,
                num_decoder_layers = self.num_decoder_layers,
                dim_feedforward = self.dim_feedforward,
                dropout = self.dropout,
                activation = self.activation,
                batch_first = self.batch_first
                )
    def forward(self, encoderSeq, decoderSeq, decoderPaddingMask=None, decoderMask=None):
        y_hat = self.transformer(src=encoderSeq, tgt=decoderSeq, tgt_mask=decoderMask, tgt_key_padding_mask=decoderPaddingMask)
        return y_hat
</code></pre>
<p>I want to get my model to produce output with part of the decoder input sequence masked out.</p>
<p>According to <a href="https://ketanhdoshi.github.io/Transformers-Arch/" rel="noopener nofollow ugc">this</a> site, if I want the last item in the decoder input sequence to be ignored in decoder self-attention, I need to generate a mask like this:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/1/1143a7531b33b938dd5c968d43eae2bd52e904ba.png" alt="Screenshot 2022-01-04 at 13.06.20" data-base62-sha1="2sJ3YKSsXb77SGtKk5htbTMIfTI" width="177" height="180"></p>
<p>Where the black squares are 0 and white squares are -inf.<br>
Code to generate the mask:</p>
<pre><code class="lang-auto">def generateDecoderMask(size, n):
    output = torch.zeros((size, size))
    if not n:
        return output
    output[:, -n:] = -torch.inf
    output[-n:, :-n] = -torch.inf
    return output
</code></pre>
<p><strong>Problematic code:</strong></p>
<pre><code class="lang-auto">model = model()
decoderMask = generateDecoderMask(model.seqLength, 1)
encSeq = torch.rand((1, model.seqLength, model.d_model))
decSeq = torch.rand((1, model.seqLength, model.d_model))
model(encSeq, decSeq, decoderMask=decoderMask)
</code></pre>
<p>The decoder mask used here is:</p>
<pre><code class="lang-auto">[[0,        0,        -inf,
  0,        0,        -inf,
  -inf,     -inf,     -inf]]
</code></pre>
<p>The output is all just nan.<br>
The problem is present on trained and untrained models.</p>
<p><strong>What I’ve tried</strong><br>
• Using a decoder mask like this:</p>
<pre><code class="lang-auto">[[0,        0,        -inf,
  0,        0,        -inf,
  0,        0,        -inf]]
</code></pre>
<p>• Using tgt_key_padding_mask instead:</p>
<pre><code class="lang-auto">decoderMask = torch.tensor([0, 0, 0, 0, 1]).bool()
model(encSeq, decSeq, decoderPaddingMask=decoderMask)
</code></pre>
<p>Both of these alternative ways produced the same output given the same input and no nan. But the mask used in the first bullet point shouldn’t be correct, should it? The last row and column represent similarity scores for padding, what’s the use of keeping the last row?</p>
<p>What am I doing wrong?</p>
<p><a href="https://pastebin.com/2kjgyuJZ" rel="noopener nofollow ugc">Code that does everything</a></p>