nn.LSTM doesn't seem to learn anything or not updating properly
I was trying out a simple LSTM use case form pytorch, with the following model.
class SimpleLSTM(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super().__init__()
        
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.lstm = nn.LSTM(batch_first=True, input_size=embedding_dim, num_layers=1, hidden_size=hidden_dim, bidirectional=True)
        self.linear = nn.Linear(hidden_dim*2, 1)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):   # NxD, padded to same length with 0s in N-sized batch
        x = self.embedding(x)
        output, (final_hidden_state, final_cell_state) = self.lstm(x)
        x = self.linear(output[:,-1,:])
        x=self.sigmoid(x)
        return x

It is a binary classification, with BCELoss (combined with the Sigmoid output layer). Unfortunately, loss is stuck at 0.6969 (i.e. it is not learning anything).
I've tried using final_hidden_state, output[:,0,:] feeding into the linear layer, but so far no dice.
Everything else (optimizer, loss criterion, train loop, val loop) already works because I tried the exact same setup with a basic NN using nn.Embedding, nn.Linear, and nn.Sigmoid only, and could get to good loss decrease and high accuracy. In the SimpleLSTM, the only thing I added is the nn.LSTM.


Typically final_hidden_state is passed to linear, not output. Use it.
add 1-2 more linear layers after the LSTM.
try lower LR, especially when embeddings are not pre-trained.
Better yet, try loading pre-trained embeddings.

