The learning rate seems to be quite high with a value of 0.1 and the 'sum' loss reduction, so I would recommend to play around with some hyperparameters.
If you get stuck, try to overfit a small dataset (e.g. just 10 samples) and scale up the use case afterwards.
Also, you are appending the loss with the entire computation graph in:
loss_train[epoch] += loss

which will increase the memory usage and you could thus easily run out of memory so detach() the tensor before the inplace accumulation.