Help understanding how to implement a LSTM


I don’t have a lot of experience with PyTorch and neural networks in general but I am trying to make a model. In this model, I want to generate the graph o a function with a set of parameters, add some noise and have the neural network find out which parameters were used in the function.

For the input, I created a function which makes the graph of a Gaussian pulse with random parameters. As for the model, I wanted to take the graph as a time sequence, pass that through an LSTM, get its last output and pass it through a linear layer for scaling. However, the model always makes the same prediction for all random sequences.

I don’t know what I’m doing wrong and need help to find out how to properly implement it. Here’s a link to Google Colab that I’m using to test the model: