Your train works, and that is a dangerous moment: now it looks right. In this step you break it in four ways — because the loop is five lines, and each one has its own mistake.

Try each by temporarily changing your train_scratch.py (and change it back). All the printed output here is real, from lr = 0.1, seed = 42.

1. You forgot zero_() — the silent failure

Delete the two lines W.grad.zero_() and b.grad.zero_(). The program does not crash. The test does not crash right away. Here are the first steps:

step   0  loss 2.3207  accuracy 10.2%
step   1  loss 2.0832  accuracy 28.5%
step   2  loss 1.7912  accuracy 51.7%
step  10  loss 0.7946  accuracy 71.4%

Indeed — faster than the correct loop (1.9216 / 58.5 % was step 2; here the loss fell even lower). But that is misleading: the gradients of earlier steps pile up, so each update is "this step's advice plus all earlier ones" — like a ball that gains speed as it rolls.

Keep looking:

step  20  loss 0.8329  accuracy 76.4%
step  30  loss 0.8907  accuracy 79.8%
step  40  loss 0.9549  accuracy 80.2%
step  50  loss 0.9814  accuracy 81.1%
step  60  loss 0.9843  accuracy 81.8%
step  61  loss nan     accuracy 9.4%
step 100  loss nan     accuracy 9.4%

This is the same mistake you saw in lesson 2 in three lines (24 → 48 → 72). Here it hides for sixty steps.

2. The update outside a no_grad block

Move W -= lr * W.grad and b -= ... out of with torch.no_grad(): to the loop level:

RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

PyTorch does not want anyone to change its tracked tensor "behind its back", because the graph would become wrong. no_grad() tells it: "you do not need to care about this, just change the number".

3. W = W - ... instead of W -= ...

with torch.no_grad():
    W = W - lr * W.grad
    b = b - lr * b.grad

AttributeError: 'NoneType' object has no attribute 'zero_'

This one is harder to connect with its cause. W - lr * W.grad is a new tensor, and the variable W now points to it; the new one is not tracked, so W.grad is None. W -= ... changes the same object, which is why that form is required.

4. You forgot requires_grad_()

W = torch.randn(784, 10) * 0.01        # without .requires_grad_()
b = torch.zeros(10)

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

loss does not remember ever being connected to W, so backward() has nowhere to go. The message says a lot: "element 0 [your loss] does not track a gradient" — so nothing that went into it was asked to be tracked.

What is left

Two of the four are loud (2, 4), one is "half" (3: the error appears far from the cause), and one is silent (1). In step 4 you saw that the test catches the silent one: loss is nan at step 61. That is its job.

When you write your own training loop next time — watch the loss over the whole run, not only at the start. Here the silent failure looked even better than the correct loop in the first steps, and only gave itself away after sixty.