🎓 A1Learner
🔍

PyTorch When Training Fails

A debugging checklist is a fixed order of checks that turns a broken training run into named faults with evidence. This lesson runs the checklist on one script with three planted faults.

The dtype error, caught

A try block can catch the error and print it as evidence. The message names both dtypes, Long and Float.

try:
    nn.Linear(3, 1)(torch.tensor([[1, 2, 3]]))
except RuntimeError as error:
    print(f"RuntimeError: {error}")

Broadcasting stretches a column

A (4, 1) tensor minus a (4,) tensor gives a (4, 4) tensor. No error appears, so the shape print is the evidence.

column = torch.tensor([[0.0], [1.0], [2.0], [3.0]])
row = torch.tensor([0.0, 1.0, 2.0, 3.0])
print((column - row).shape)

A sane rate on the same data

The same regression with a rate of 0.05 falls from the first steps. The fault was the rate, not the data or the model.

small = nn.Linear(1, 1)
with torch.no_grad():
    small.weight.fill_(0.0)
    small.bias.fill_(0.0)
small_opt = torch.optim.SGD(small.parameters(), lr=0.05)
for step in range(3):
    small_opt.zero_grad()
    small_loss = loss_fn(small(x), y)
    small_loss.backward()
    small_opt.step()
    print(f"Step {step}: loss {float(small_loss):.4f}")

Example

import torch
from torch import nn

torch.manual_seed(0)
x = torch.linspace(-1, 1, 20).view(-1, 1)
y = 2 * x + 1
loss_fn = nn.MSELoss()

print("Fault 1: a learning rate of 5.0 on a small regression")
model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(0.0)
    model.bias.fill_(0.0)
optimizer = torch.optim.SGD(model.parameters(), lr=5.0)
for step in range(5):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
    print(f"Step {step}: loss {float(loss):.2f}")

print("Fault 2: long inputs fed to a float model")
layer = nn.Linear(3, 1)
long_input = torch.tensor([[1, 2, 3], [4, 5, 6]])
try:
    layer(long_input)
    print("No error raised")
except RuntimeError as error:
    print(f"RuntimeError: {error}")

print("Fault 3: predictions (16, 1) against targets (16,)")
predictions = torch.linspace(0, 1, 16).view(-1, 1)
targets = torch.arange(16).float()
difference = predictions - targets
print(f"Subtraction shape: {tuple(difference.shape)}")
print(f"Silent loss: {float(loss_fn(predictions, targets)):.4f}")
print(f"Correct loss: {float(loss_fn(predictions, targets.view(-1, 1))):.4f}")

print("The repaired script: rate 0.05, float inputs, targets shaped (20, 1)")
model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(0.0)
    model.bias.fill_(0.0)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
for step in range(61):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
    if step % 15 == 0 or step == 60:
        print(f"Step {step}: loss {float(loss):.4f}")
print(f"Learned weight: {float(model.weight):.3f}, bias: {float(model.bias):.3f}")

Key points

  • Check the loss trace first, because an exploding trace names a rate fault before any code reading
  • Check dtypes second, because a dtype fault stops the run with an error that names both types
  • Check shapes third, because a broadcasting fault raises no error and only a shape print exposes it
  • Broadcasting: a (16, 1) tensor minus a (16,) tensor silently gives a (16, 16) tensor
  • Each fault is fixed by one change, and the repaired run proves the fix by falling cleanly
  • The repaired script changes the rate, feeds float inputs, and shapes the targets as (20, 1)

In this example, one small regression script carries three planted faults, and the checklist finds each by evidence. Fault 1 is the rate. At a learning rate of 5.0 the loss explodes from 2.47 to over 43 million in five steps. Fault 2 is the dtype. Long inputs fed to a float Linear layer raise a RuntimeError that names Long and Float. Fault 3 is silent broadcasting. Predictions shaped (16, 1) against targets shaped (16,) subtract into a (16, 16) tensor. The wrong loss of 70.3444 prints beside the correct 67.5111. The repaired script then runs the same data at a rate of 0.05. Its loss falls from 2.4737 to 0.0163, and the learned weight and bias land near the true 2 and 1.

Note

The silent fault is the dangerous one, since its loss still prints as a plausible number.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Which fault is found only by printing a shape, and why?