Code reviewMediumcr-g628

Subject Model training pytorchLevel Mid–Senior~18 minCommon in Code quality & review interviewsIndustries Software development

Question

Review this PyTorch training loop.

Training is unstable and the loss behaves erratically. What's missing?

What a strong answer looks like

Separate real bugs from style. Rank issues by severity, point at the root cause rather than the symptom, and suggest a concrete fix — specific and kind.

Learn the concepts

Talk through your review

Code to reviewpython

1import torch2import torch.nn as nn3 4def train(model, loader, epochs=10, lr=1e-3):5    opt = torch.optim.Adam(model.parameters(), lr=lr)6    loss_fn = nn.CrossEntropyLoss()7    for epoch in range(epochs):8        for xb, yb in loader:9            logits = model(xb)10            loss = loss_fn(logits, yb)11            loss.backward()12            opt.step()13        print(f"epoch {epoch} loss {loss.item():.4f}")

Run or narrate your approach, then ask the coach.