Test PyTorch models like you mean it. pytest-torch is a pytest plugin with small, composable helpers for the tests every model needs, and failure messages that say where a mismatch is and how large it is.

Status: alpha. The API may still change. Requires Python 3.9+, pytest 7+ and PyTorch 2.0+. Documentation: https://akj7.github.io/pytest-torch/

pip install pytest-torch

Start here

Pick the question you're asking. Each path is one or two calls; every snippet is runnable in examples/test_start_here.py.

1. Check one prediction

Does the model behave like an inference model, and do its outputs look right?

from pytest_torch import Probabilities, TensorSpec, assert_output_spec, check_inference_contract

SPEC = {"logits": TensorSpec(("B", 3)), "probs": Probabilities(("B", 3))}

def test_one_prediction(model, make_tensor):
    x = make_tensor(8, 16)
    check_inference_contract(model, x)       # eval mode, finite, deterministic, batch-invariant, devices
    with torch.inference_mode():
        assert_output_spec(model(x), SPEC, B=8)   # shapes, dtypes, ranges, probabilities

model is your fixture; make_tensor and torch_device come with the plugin, so the test runs on every available device with seeded inputs. Go further: invariance (padding, masks, permutations), output specs, unusual inputs and dtypes.

2. Check a model change

Does the refactored, ported, compiled or compressed model still match the original?

from pytest_torch import assert_matches_reference, assert_models_agree, localize_divergence

def test_refactor_matches_the_original(original, refactored, make_tensor):
    x = make_tensor(8, 16)
    assert_matches_reference(refactored, original, x)              # same outputs, path by path
    localize_divergence(original, refactored, x).raise_if_diverged()  # ...and layer by layer

def test_bf16_model_stays_close(original, make_tensor):
    x = make_tensor(256, 16)
    bf16 = copy.deepcopy(original).to(torch.bfloat16)
    assert_models_agree(original, bf16, x, inputs_candidate=x.to(torch.bfloat16), min_top1=0.95)

Go further: deployment parity (torch.compile, torch.export, AOTInductor, ONNX, autocast), streaming and stateful models, tolerances.

3. Check a release

Is the model we're about to ship as good, as compatible and as fast as the last one?

from pytest_torch import assert_checkpoint_loads

def accuracy(output, batch):
    return (output["logits"].argmax(-1) == batch[1]).float()   # per sample: paired comparisons

def test_release_accuracy(model, val_loader, metric_baseline):
    metric_baseline(model, val_loader, {"accuracy": accuracy}, minimum={"accuracy": 0.9},
                    context={"dataset": "synthetic", "split": "validation"})

def test_last_release_checkpoint_loads():
    assert_checkpoint_loads(Classifier, "tests/data/classifier_v1.pt")

def test_speed(model, make_tensor, perf_baseline):
    perf_baseline(model, make_tensor(32, 16))

Record baselines once with --torch-update-metrics and --torch-update-perf, commit them, and CI compares against them. Go further: dataset metrics (paired comparisons with confidence intervals), checkpoints, performance, failure reproduction.

Or generate the tests

ModelSuite turns one class into a test suite: pick the checks, and each becomes a test on every device.

from pytest_torch import ModelSuite

class TestClassifier(ModelSuite):
    checks = ("contract", "compile", "export", "output_spec", "checkpoint", "training")
    output_spec = {"logits": TensorSpec(("B", 3))}

    def make_model(self):
        return Classifier()

    def make_inputs(self, make_tensor):
        return make_tensor(8, 16)

    def loss(self, output, inputs):          # only needed for "training"
        return output["logits"].logsumexp(-1).mean()

See generated test suites for every check and option. The helpers it uses stay available on their own.

Everything else

Installation

pip install pytest-torch          # or: uv add --dev pytest-torch

PyTorch is not installed for you (pick the build that matches your hardware). The plugin registers itself through pytest's pytest11 entry point, and importing it never imports torch. Optional extras:

Examples

Runnable test files in examples/ (pytest examples/ -v). Most have a "Try ..." comment: a one-line bug that makes the test fail.

examples/integration/ has three small but realistic projects, laid out like real repositories and tested end to end: a transformer text classifier, an image classifier with its deployment pipeline, and a small GPT with a KV cache.

Development

uv sync                          # environment, lockfile, dev dependencies
uv run pytest                    # the plugin's own tests
uv run pytest examples/ -v       # runnable examples of every feature
uv run pytest -m "not slow"      # skip AOTInductor compilation and other slow tests
uv run --group docs mkdocs serve # documentation at http://127.0.0.1:8000
uv build                         # sdist and wheel

CI (.github/workflows/ci.yml) runs the tests and the examples over operating systems, Python versions and torch versions with CPU-only wheels, plus the slow tests once on the newest torch. docs.yml builds the docs on every PR and publishes them from main; release.yml publishes to PyPI from a GitHub release.

License

MIT