The loss fell to 11.25 and the model learned nothing

Broadcasting is a feature that is willing to make sense of shape combinations you did not intend to write. In the chapter’s opening program the targets arrive as a column, the predictions as a row, and the subtraction silently produces a four by four grid of errors: every prediction compared against every target. Nothing raises, the loss falls by two thirds, and the model ends up answering y = 7.5 whatever the input.

This program composes the chapter’s pieces into one run: the opening training loop, the shape diagnosis that names the broadcast, and the fixed point the broadcast implies. It then runs the repaired loop, so the repair is judged by the quantity the diagnosis said would change, not by the loss going down.

Run it

Save this as tensor_broadcast.py and run it with Python 3.11 or later. It needs PyTorch; install it with pip install torch, then run python tensor_broadcast.py.

"""The tensor flowing through the loop: a (4,) prediction against a (4, 1) target."""
import torch


def train(x, target, steps=500, lr=0.05):
    w = torch.tensor(1.0, requires_grad=True)
    b = torch.tensor(0.0, requires_grad=True)
    for _ in range(steps):
        prediction = w * x + b
        loss = ((prediction - target) ** 2).mean()
        loss.backward()
        with torch.no_grad():
            w -= lr * w.grad
            b -= lr * b.grad
            w.grad.zero_()
            b.grad.zero_()
    return w.item(), b.item(), loss.item()


x = torch.tensor([1.0, 2.0, 3.0, 4.0])
target = torch.tensor([[3.0], [6.0], [9.0], [12.0]])   # a column, as a loader emits

w, b, loss = train(x, target)
prediction = w * x + b
errors = prediction - target

print(f"prediction {tuple(prediction.shape)}  target {tuple(target.shape)}")
print(f"error matrix {tuple(errors.shape)}: {errors.numel()} errors for 4 examples")
print(f"converged  w={w:.4f}  b={b:.4f}  loss={loss:.6f}")
print(f"prediction at x=10: {w * 10 + b:.4f}")

flat = target.flatten()
print(f"mean of targets {flat.mean().item():.4f}   population variance {flat.var(unbiased=False).item():.4f}")

w2, b2, loss2 = train(x, target.reshape(-1))
print(f"repaired  w={w2:.4f}  b={b2:.4f}  loss={loss2:.6f}")
print(f"prediction at x=10: {w2 * 10 + b2:.4f}")

The run prints the shapes that disagree, the converged parameters, and the predicted numbers.

Expected outcome

prediction (4,)  target (4, 1)
error matrix (4, 4): 16 errors for 4 examples
converged  w=0.0013  b=7.4963  loss=11.250002
prediction at x=10: 7.5090
mean of targets 7.5000   population variance 11.2500
repaired  w=2.9999  b=0.0003  loss=0.000000
prediction at x=10: 29.9992

The error is in the shape, not in the loop

Align the two shapes from the right:

prediction       (4)
target       (4, 1)

The rightmost pair is 4 against 1, so the 1 expands, and the missing leading dimension of prediction expands with it. The difference is (4, 4), and .mean() averages sixteen values rather than four. The objective is not the error of each prediction against its own target, it is the error of every prediction against every target.

That objective has an analysable minimum: the prediction that minimises the average of (p - t_j)^2 over all targets is the mean of the targets, and the residual loss there is the mean squared deviation from it. The run prints both numbers before the repair: the mean is 7.5000, the variance 11.2500. The broken loop converges to b = 7.4963 with a loss of 11.250002. The diagnosis came first, and the numbers confirmed it.

The repair is to state the contract rather than to reshape until the numbers change:

assert prediction.shape == target.shape, (
    f"prediction {tuple(prediction.shape)} vs target {tuple(target.shape)}"
)

Placed just before the loss, that assertion fails immediately, in the loop body, naming both shapes. Whether you flatten the target or give the prediction a trailing axis depends on what the rest of the program expects. What matters is that you decided.

Mechanism and limitations

Documented: the chapter shows the opening loop converging to b = 7.4963 with w = 0.0013 and a loss of 11.250002, and accounts for both numbers: the mean of the targets is 7.5, the mean squared deviation from it is 11.25.

Observed: the run above. The prediction was (4,) and the target (4, 1), so the error matrix was (4, 4): sixteen errors for four examples. The loop converged to w = 0.0013, b = 7.4963 and loss 11.250002, and answered 7.5090 at x = 10. The mean and variance printed as 7.5000 and 11.2500. After flattening the target, the same loop reached w = 2.9999, b = 0.0003 and answered 29.9992 at x = 10.

Not shown. The program composes the chapter’s fragments into one run and adds the shape prints and the repaired loop, so it is not a copy of any single listing. It tests one shape, one batch size and one seed. A falling loss is not evidence here: the evidence is the element count, the converged parameters and the prediction at x = 10. Nothing here shows that a shape assertion can detect two equal-sized axes being swapped, nor what put the column there.

Canonical source , the chapter this example composes. Read the chapter .