You write nothing here. Below is a real Python console session: we look at what the new parts do and check, number by number, that it is the same thing you did by hand in lesson 4.

Same data: the first 10 000 FashionMNIST images (load_fashion("train", n=10_000) from fashion.py), with torch.manual_seed(42) first.

1. Batches

>>> loader = DataLoader(TensorDataset(x, y), batch_size=64, shuffle=True)
>>> len(loader)
157
>>> xb, yb = next(iter(loader))
>>> xb.shape, yb.shape
(torch.Size([64, 784]), torch.Size([64]))
>>> yb[:8]
tensor([9, 4, 9, 8, 5, 4, 2, 7])
>>> [len(b[0]) for b in loader][-2:]
[64, 16]

TensorDataset(x, y) pairs each image with its label and DataLoader deals them out in batches. shuffle=True means a different order every epoch.

2. The model as one object

>>> model = nn.Linear(784, 10)
>>> [p.shape for p in model.parameters()]
[torch.Size([10, 784]), torch.Size([10])]
>>> sum(p.numel() for p in model.parameters())
7850

3. One update — two routes, the same result

>>> loss_fn = nn.CrossEntropyLoss()
>>> loss = loss_fn(model(xb), yb)
>>> loss.item()
2.358222723007202
>>> loss.backward()
>>> w_before = model.weight.detach().clone()
>>> grad = model.weight.grad.clone()
>>> by_hand = w_before - 0.1 * grad               # lesson 4's formula
>>> torch.optim.SGD(model.parameters(), lr=0.1).step()

The story of one weight ([0, 300]):

before           -0.012392
gradient         -0.033848
after (by hand)  -0.009007     -0.012392 - 0.1 * (-0.033848)
after (optimizer)-0.009007

>>> torch.allclose(by_hand, model.weight.detach())
True
>>> (by_hand - model.weight.detach()).abs().max().item()
3.725290298461914e-09
>>> loss_fn(model(xb), yb).item()      # same batch, after one update
2.0123372077941895

What to notice

1. A batch is the same kind of matrix as in lesson 4 — just 64 rows. xb is [64, 784]. Your whole formula works unchanged. 10 000 images = 156 full batches of 64 plus one leftover of 16 (156·64 + 16 = 10000), so one epoch = 157 updates, not one.

2. nn.Linear holds the same 7 850 numbers, stored transposed. In lesson 4 W was [784, 10]; nn.Linear(784, 10) stores weight as [10, 784] and computes x @ weight.T + bias itself. Same count: 784·10 + 10. model.parameters() lists all of them — which is why the optimizer no longer needs you to tell it what to update.

3. optimizer.step() is not magic. It is exactly W -= lr * W.grad: the difference between "by hand" and "optimizer" is 3.7·10⁻⁹, which is just floating-point rounding. And the loss on that same batch dropped from 2.358 to 2.012 after one update.

The numbers here are from my run. If you repeat exactly the same sequence you get the same ones; otherwise different batches and different starting weights give different numbers — but you will still see the three observations.