tom:

good to see you around!


Thank you, Thomas! You too!


So drop the .data , wrap in no_grad works here as far as I can see:
def perfect_update(x):
    with torch.no_grad():
        x -= 0.0001 * x.grad
    x.grad.zero_()
    return x

(I’m having reservations about calling it update when it also zeros grad and about returning x, too, but hey.)


Oh, I see, this way wasn’t working since I was trying the direct: x = x - lr*grad approach:
    with torch.no_grad():
        x = x - 0.0001 * x.grad
    _= x.grad.zero_()

which fails: AttributeError: 'NoneType' object has no attribute 'zero_' in the x.grad.zero_() call, so I guess in this case in-place subtraction is what I was missing. Then I switched to in-place, and forgot to try this approach, and other ways were failing.
So yes, your perfect_update does the trick!
In-place modification of x is a must then.
Thank you, Thomas.

(I’m having reservations about calling it update when it also zeros grad and about returning x, too, but hey.)

Well, that was just an abstraction I added to make my question more readable with the different x update versions. You’re right, it’s not an update.
the full clean solution now goes:
# given A and b, and Ax=b: find x using SGD
A = torch.randn(6, 6, requires_grad=False)
b = torch.randn(6, 1, requires_grad=False)

# choose random x and then search with SGD to find x-hat that is the closest to x in Ax=b
x = torch.randn(6, 1, requires_grad=True)

for i in range(40000):
    loss = torch.norm(A @ x - b)
    loss.backward()
    with torch.no_grad():
        x = x - 0.0001 * x.grad
    _= x.grad.zero_()
    if not i % 5000: print(f"{loss: >.6f}    {x.T}")

Excellent!