LSTM layer outputing nan on second batch

petartushev: <p>I have an email dataset which I have to classify as spam or not. I am using <code>gensim.models.Word2Vec</code> for creating word embeddings of the train and test emails. For the processing of the dataset, I am using the tokens</p>
<pre><code class="lang-auto">MAX_SENTENCE_TOKENS = 100
EMBEDDING_SIZE = 300
START_TOKEN = '&lt;START&gt;'
END_TOKEN = '&lt;END&gt;'
PADDING_TOKEN = '&lt;PADDING&gt;'
</code></pre>
<p>for creating sentances that are at maximum 100 tokens long. After this, I train a skip-gram Word2Vec model and I exchange the actual words for the embedding indices. For the dataset I have the following <code>torch.utils.data.Dataset</code> class defined:</p>
<pre><code class="lang-auto">class Data(Dataset):
    def __init__(self, x, y):
        
        self.X = torch.Tensor(x, device=device).long()
        
        self.y = torch.Tensor(y, device=device).long()
        
        self.len = self.X.shape[0]
       
    def __getitem__(self, index):
        return self.X[index], self.y[index]
   
    def __len__(self):
        return self.len
</code></pre>
<p>Moreover, my hyperparameters and model are defined as:</p>
<pre><code class="lang-auto">BATCH_SIZE = 8
EPOCHS = 30
HIDDEN_SIZE = 64
NUM_LAYERS = 16
LEARNING_RATE = .0001
BIDIRECTIONAL = False

loss_fn = nn.BCELoss()
optim = Adam(model.parameters(), lr=LEARNING_RATE)

class LSTMSpamClassifier(nn.Module):
    
    def __init__(self, hidden_dim, num_layers):
        
        super(LSTMSpamClassifier, self).__init__()
        
        self.emb = Embedding.from_pretrained(torch.Tensor(w2v.wv.vectors))
        
        self.lstm1 = LSTM(input_size=EMBEDDING_SIZE, hidden_size=HIDDEN_SIZE, num_layers=NUM_LAYERS, bidirectional=BIDIRECTIONAL)
        
        self.lstm2 = LSTM(input_size=HIDDEN_SIZE, hidden_size=HIDDEN_SIZE, num_layers=NUM_LAYERS, bidirectional=BIDIRECTIONAL)
        
        self.lstm3 = LSTM(input_size=HIDDEN_SIZE, hidden_size=HIDDEN_SIZE, num_layers=NUM_LAYERS, bidirectional=BIDIRECTIONAL)
        
        self.l1 = Linear(HIDDEN_SIZE, 128)
        
        self.l2 = Linear(128, 16)
        
        self.l3 = Linear(16, 1)

        
    def forward(self, x):

        embedding = self.emb(x)

        out, states = self.lstm1(embedding)

        out, states = self.lstm2(out)   # &lt;-- the output of this lstm layer is nan
        
        out, states = self.lstm3(out)
        
        final_hidden_state = out[:, -1, :]

        x = F.relu(self.l1(final_hidden_state))
        
        x = F.relu(self.l2(x))
        
        out = F.sigmoid(self.l3(x))
        
        return out
</code></pre>
<p>With my training loop being:</p>
<pre><code class="lang-auto">model.train()

for epoch in range(EPOCHS):
    
    for x, y in train_dataloader:
        
        y_pred = model(x)

        y = y.float()
        y_pred = y_pred.squeeze(-1)

        loss = loss_fn(y, y_pred)
        
        optim.zero_grad()

        loss.backward()
        
        optim.step()
        
        
    print(f'EPOCH: {epoch}/{EPOCHS} | Loss: {loss}')

</code></pre>
<p>When I try to train this model the first batch is as I would expect, but the second batch on the second lstm layer, it outputs a <code>torch.Tensor</code> of nans.</p>
<p>Example output:</p>
<pre><code class="lang-auto">tensor(51.2606, grad_fn=&lt;BinaryCrossEntropyBackward0&gt;)
EPOCH: 0/30 | Loss: 51.260623931884766
tensor(nan, grad_fn=&lt;BinaryCrossEntropyBackward0&gt;)
EPOCH: 1/30 | Loss: nan

</code></pre>
<p>Does anyone have an explanation of this behaviour?</p>

vdw: <p>Just some comments and questions</p>
<ul>
<li>
<p>Why do you have <code>lstm1</code>, <code>lstm2</code>, and <code>lstm3</code>, with each of this layers having 16 layers. This means you have a LSTM layer with <code>3*16=48</code> layers. A single <code>lstm</code> with maybe 1-3 layers should be more than enough.</p>
</li>
<li>
<p>I assume the shape of <code>embedding</code> is <code>(batch_size, embedding_size)</code>. However, you define the LSTM layers with <code>batch_first=False</code> (that’s the default). So you probably need to do <code>embedding = embedding.transpose(1,0)</code> first.</p>
</li>
<li>
<p>With <code>batch_first=False</code>, the shape of <code>out</code> is <code>(seq_len, batch_size, 2*hidden_size)</code>. So <code>out[:, -1, :]</code> is wrong, and it should be simply <code>out[-1]</code></p>
</li>
</ul>

petartushev: <ul>
<li>I have (or had now) so many layers in <code>lstm</code> because I thought I needed that many.</li>
<li>Instead of transposing he embeddings and re-selecting the output of the final <code>lstm</code> I added <code>batch_first=True</code> to all of the <code>lstm</code> layers. Still, the output is:</li>
</ul>
<pre><code class="lang-auto">tensor(51.3910, grad_fn=&lt;BinaryCrossEntropyBackward0&gt;)
EPOCH: 0/30 | Loss: 51.39103698730469
tensor(nan, grad_fn=&lt;BinaryCrossEntropyBackward0&gt;)
EPOCH: 1/30 | Loss: nan
</code></pre>