NaN values when I pass key_padding_mask to MultiHeadAttention

rajeshTU: <p>My input data looks like</p>
<pre><code class="lang-auto">tensor([[[ 18.6752, -11.8400,   5.1000,   1.4900,  -3.6898,   0.5981,   1.6500,
            2.0000],
         [ 20.6315,   3.3100,   3.5000,   1.8200, -18.5913,  -0.1733,  -3.1300,
            2.0000],
         [ 20.8761,   6.5100,  10.2500,   2.4100,  -4.9920,   0.0291,   3.1200,
            1.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,
            0.0000]]])
</code></pre>
<p>where  [  0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,<br>
0.0000] is padding.</p>
<p>My forward pass looks like this</p>
<pre><code class="lang-auto">x = sensor_data[i].to('cuda')
mask = (x == 0)  # Create padding mask
mask = mask.any(dim=2)
mask = torch.transpose(mask, 0, 1)
x = self.posEmbed(x)
x = self.inputEmbed(x)
encoded_sensor = self.encoder(x, mask)
</code></pre>
<p>And the encoder layer looks like this</p>
<pre><code class="lang-auto">class EncoderLayer(nn.Module):
    def __init__(self, embed_dim, num_heads, feedforward_dim, dropout_prob=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout_prob)
        self.layer_norm1 = nn.LayerNorm(embed_dim)
        self.feedforward = nn.Sequential(
            nn.Linear(embed_dim, feedforward_dim),
            nn.ReLU(),
            nn.Linear(feedforward_dim, embed_dim),
        )
        self.layer_norm2 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout_prob)

    def forward(self, x, key_padding_mask):
        # x: (seq_length, batch_size, embed_dim)
        # self-attention
        attn_output, _ = self.self_attention(x, x, x, key_padding_mask=key_padding_mask)
        attn_output = self.dropout(attn_output)
        x = self.layer_norm1(x + attn_output)

        # feed-forward
        ff_output = self.feedforward(x)
        ff_output = self.dropout(ff_output)
        x = self.layer_norm2(x + ff_output)

        return x
</code></pre>
<p>After the forward pass through encoder, the places with padding becomes nan. But when I remove the key_padding_mask argument it works well. Please help me with this asap. thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Lars_Hertel: <p>I have the same issue. Did you figure out a solution <a class="mention" href="/u/rajeshtu">@rajeshTU</a> ?</p>