Hi Alger_none!
I’ve only had a really quick look at your code - I can have a more in depth look if this suggestion doesn’t work.
Specifically, I’m looking at your train function. It looks as though you’re trying to zero_grad() your model - not your optimizer. My suggestion would be to refactor your train function have:

zero_grad() your optimizer, to avoid accumulating gradients. Usually do this at the start of your training loop.

step() your optimizer after `loss.backward()

This tutorial 2 should help you out!