Torch-TensorRT compiles a PyTorch model into an optimized TensorRT engine with no separate export or runtime. This example quantizes a PyTorch / HuggingFace model with NVIDIA Model Optimizer and then compiles the quantized graph in-framework with Torch-TensorRT for deployment.

Quantization is an effective model optimization technique that compresses your models. Model Optimizer inserts Q/DQ nodes into the eager PyTorch graph; torch_tensorrt.compile(ir="dynamo") then converts those Q/DQ nodes into native TensorRT FP8 precision layers, following the Torch-TensorRT quantization guide.

This section focuses on the in-framework Torch-TensorRT path: a PyTorch front end (mtq.quantize) feeding a Dynamo-compiled TensorRT engine, demonstrated end-to-end on a HuggingFace ViT image classifier. If you instead want a portable ONNX → TensorRT artifact, or you start from an ONNX model, see the sibling torch_onnx and onnx_ptq examples (compared in the Support Matrix).

Please use the TensorRT docker image (e.g., nvcr.io/nvidia/tensorrt:26.02-py3) or visit our installation docs for more information.

docker run --gpus all -it --rm -v $(pwd):/workspace -w /workspace nvcr.io/nvidia/tensorrt:26.02-py3 bash

Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies.

pip install -U "nvidia-modelopt[hf]"
pip install -r requirements.txt

The low-precision kernels Torch-TensorRT emits need a GPU that supports the target format:

Quantize a HuggingFace ViT, then compile the Q/DQ graph with Torch-TensorRT into a single torch.nn.Module you call from PyTorch:

import torch
import torch_tensorrt

import modelopt.torch.quantization as mtq
from modelopt.recipe import load_recipe
from modelopt.torch.quantization.utils import export_torch_mode

# 1. Quantize the eager PyTorch model with a Model Optimizer PTQ recipe.
recipe = load_recipe("model_type/vit/ptq/fp8")
mtq.quantize(model, recipe.quantize.model_dump(), forward_loop=calibrate)

# 2. Compile the quantized (Q/DQ) graph with Torch-TensorRT.
#    export_torch_mode() makes Model Optimizer emit Q/DQ in the TRT-friendly form,
#    and min_block_size=1 lets single-node Q/DQ + matmul subgraphs become TRT
#    precision layers (per the Torch-TensorRT quantization guide).
with export_torch_mode():
    trt_model = torch_tensorrt.compile(
        model,
        ir="dynamo",
        min_block_size=1,
        truncate_double=True,
        inputs=[torch_tensorrt.Input(
            min_shape=(1, 3, 224, 224),
            opt_shape=(128, 3, 224, 224),
            max_shape=(1024, 3, 224, 224),
            dtype=torch.float16,
        )],
    )

logits = trt_model(pixel_values)  # call it like any nn.Module

The runnable script torch_tensorrt_ptq.py wraps this flow end-to-end. It:

1. Loads a HuggingFace ViT classifier (default google/vit-large-patch16-224).
2. Builds a tiny calibration loader from zh-plus/tiny-imagenet (avoids the gated ILSVRC/imagenet-1k repo, so the example runs unauthenticated).
3. Runs mtq.quantize with one of the recipes under modelopt_recipes/ (see ViT Recipes).
4. Saves the quantized Model Optimizer state (FP16 weights + Q/DQ metadata) to <save_dir>/vit_modelopt_state.pt for reuse without recalibration (see Custom Recipes).
5. Compiles the quantized model with torch_tensorrt.compile and verifies that the compiled-model argmax matches the fake-quant argmax on a sample input.

# Default model is google/vit-large-patch16-224, default recipe is the ViT FP8 recipe.
python torch_tensorrt_ptq.py --calib_samples 1024 --batch_size 128

# Quantize but don't TRT-compile (handy on a non-TRT host).
python torch_tensorrt_ptq.py --skip_trt

All three of these examples reach the same destination — a low-precision TensorRT engine — but quantize at a different point in the pipeline and emit a different artifact, so they suit different deployment stacks:

This example and torch_onnx share the same PyTorch front end (mtq.quantize), so the numerics are identical — they differ only in the back end: this one keeps the graph in-process and hands it to Torch-TensorRT, while torch_onnx exports a portable ONNX artifact for the standalone TensorRT runtime. onnx_ptq instead quantizes the ONNX graph directly, for when you start from an ONNX model rather than PyTorch. Pick this example when your serving stack is PyTorch-native and you'd rather avoid an ONNX export step.

This is the recipe the CLI selects by default when --model_id points at a HF ViT classifier. It is tuned for the HF ViT module layout and is composed from the shared $import building blocks under modelopt_recipes/configs/ (ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8}) rather than spelling out each quant_cfg entry.

Script — quantize and (optionally) Torch-TensorRT-compile a ViT.

# Custom model + custom recipe, saving the quantized state elsewhere.
python torch_tensorrt_ptq.py \
    --model_id <huggingface/model-id> \
    --recipe <recipe-path-relative-to-modelopt_recipes-or-absolute-yaml> \
    --save_dir ./my_quantized

# Dump the compiled engine's per-layer info to inspect FP8 fusion.
python torch_tensorrt_ptq.py --layer_info_path ./vit_fp8_layers.txt

Script — quantize, compile, and score on ImageNet (see Evaluate Accuracy).

torch_tensorrt_accuracy.py reuses the quantize → compile pipeline above and reports ImageNet-1k top-1 / top-5 accuracy via the onnx_ptq example's evaluate() harness (examples/onnx_ptq/evaluation.py):

python torch_tensorrt_accuracy.py \
    --recipe model_type/vit/ptq/fp8 \
    --batch_size 128 \
    --baseline \
    --eval_data_size 5000 \
    --results_path results.csv

• --baseline also scores the unquantized model. It is Torch-TensorRT-compiled the same way as the quantized model, so every reported number comes from the same TRT runtime (pass --skip_trt to score the eager / fake-quant models instead).
• The eval uses a dynamic engine (default --batch_size 128) for both precisions, so it serves the trailing partial batch at any batch size.
• --results_path results.csv writes the metrics table (Metric, Top1 (%), Top5 (%)) to CSV.

Use --recipe <path> to plug in a different recipe — either a path relative to modelopt_recipes/ (resolved against the built-in recipe library) or an absolute filesystem path to a YAML file. The recipe is loaded via modelopt.recipe.load_recipe, must declare metadata.recipe_type: ptq and a quantize: section, and its quantize config is passed straight to mtq.quantize. See the existing modelopt_recipes/model_type/vit/ptq/*.yaml for the patterns used here.

torch_tensorrt_ptq.py always saves the quantized Model Optimizer state to <save_dir>/vit_modelopt_state.pt (default --save_dir ./modelopt_quantized) via mto.save. To reload it without recalibrating, restore it onto a freshly-loaded model before the TRT compile step:

import modelopt.torch.opt as mto

mto.restore(model, "./modelopt_quantized/vit_modelopt_state.pt")

• 📅 Roadmap
• 📖 Documentation
• 🎯 Benchmarks
• 💡 Release Notes
• 🐛 File a bug
• ✨ File a Feature Request