I was playing around with your code, and using exactly your parameters, but including momentum (0.9) and weight_decay(1e-5), you get a better fit. However, it still takes 1000’s of epochs to get there, way more than the 50 mentioned in the paper. The best looking fit I found was using,

4 layers of 512 neurons each
max_epochs = 2000
lr = 1e-2

By 800 epochs, you have almost a perfect fit. However, the loss seems a bit iffy. I will still look into it. I’m curious as to how they managed to get that result in 50 epochs because that is blowing my mind.
Edit: Wanted to mention that this worked only for a frequency of 2. When changed to 4, it cannot estimate the full wave.