jpeg729:

You need to keep some of the “hidden state” of the LSTM, not of the “cell state”, right?


Yes, indeed.



 jpeg729:

input_lstm is of shape (timesteps, features) right?


Yes again.
After reading your link, I am wondering how smth’s answer from here hidden states for all timesteps 5 fits in. I already incorporated this suggestion and it seems to work well. Thanks



 jpeg729:

On thing that seems weird to me is that you are backpropagating from only part of the output of the LSTM, in which case many of the parameters of the LSTM won’t get gradients at all, and others may be over/under-used.


I see your point regarding the addition of a Linear Layer to be able to distribute the gradient updates in a better way. I will certainly try this tomorrow. Do you think this might be the cause of the vanishing gradients?
One more clue I’d like to add here is that decreasing the learning rate actually causes a much slower vanishing of gradients. .