Solved by mariosasko in post #2 
              
              
                Replace the line loss_fn.backward(x) with loss_fn.backward(). After that change the optimization should behave as expected. To understand the difference between these two calls, please check the autograd tutorial (it contains both examples). 

This statement does not hold. Derivative of x**2 is 2*x …