I Built a Diagnostic Toolkit for PyTorch Because I Was Tired of Guessing Why Models Fail

Every time a PyTorch model refuses to learn, the debugging process looks the same:

After 17 years in distributed systems and SRE, I know this pattern — it is monitoring by vibes. In production infrastructure, we would never accept "the service seems slow" as a diagnostic. We measure. We trace. We verify.

So I built torchdiag — five diagnostic commands that answer the actual questions.

pip install torchdiag

PyTorch model health diagnostics — built from an SRE perspective.

Stop guessing why your model isn't learning. torchdiag

gives you five diagnostic commands that answer the questions that matter: Are gradients flowing? Are neurons alive? Did the optimizer actually update weights?

pip install torchdiag
python
import torch
import torch.nn as nn
import torchdiag
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 64),
    nn.ReLU(),
    nn.Linear(64, 10),
)

torchdiag.summary(model)

x = torch.randn(100, 784)
torchdiag.check_dead_neurons(model, x)

torchdiag.verify_step(
    model,
    torch.optim.Adam

…

import torchdiag
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 64),
    nn.ReLU(),
    nn.Linear(64, 10),
)

torchdiag.summary(model)

Prints parameter count per layer, total/trainable/frozen breakdown, memory footprint, device placement, and dtype distribution. Flags frozen parameters, split-device models, and dtype mismatches.

loss = nn.CrossEntropyLoss()(model(x), target)
loss.backward()

torchdiag.check_gradients(model)

Reports gradient mean, max, and min per layer. Flags vanishing gradients (max below 1e-7), exploding gradients (max above 100), and disconnected parameters (None gradients).

x = torch.randn(100, 784)
torchdiag.check_dead_neurons(model, x)

A dead ReLU neuron outputs zero for every input. Its gradient is permanently zero. It will never learn again. This command tells you how many you have and where. Flags critical layers with more than 50% dead neurons.

torchdiag.verify_step(
    model,
    torch.optim.Adam(model.parameters()),
    nn.CrossEntropyLoss(),
    torch.randn(32, 784),
    torch.randint(0, 10, (32,)),
)

Runs one complete training step — forward, loss, backward, optimizer step — and verifies each stage works. Confirms output shape is correct, loss is finite, gradients are computed, and parameters actually change.

Run this before your training loop. If something is broken, you will know in 1 step instead of 100 epochs.

torchdiag.memory_report()

Reports CPU RSS, GPU allocated/cached/peak per device, and MPS memory on Apple Silicon. Flags when GPU utilization exceeds 90%.

I spent 11 years at VMware working on distributed systems observability. The first thing you learn in SRE: never trust a system you cannot measure.

PyTorch models are systems. They have inputs, internal state, and outputs. When they fail, they fail silently — the loss just stays flat. No error. No exception. Just a number that does not move.

torchdiag makes the internal state visible. Five commands. No configuration. No dependencies beyond PyTorch.

PyPI: pypi.org/project/torchdiag

GitHub: github.com/AddyM/torchdiag

CI: Tests pass across Python 3.9 to 3.12

Contributions welcome. If you have a debugging pattern you use repeatedly, open an issue — it probably belongs in the toolkit.