Add --trt_build flag to torch_quant_to_onnx and simplify tests

Move TRT engine build logic into the script as a --trt_build flag,
removing the duplicate trtexec invocation from the test file.

Signed-off-by: ajrasane <arasane@nvidia.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>

Author: ajrasane

examples/torch_onnx/torch_quant_to_onnx.py

@@ -17,6 +17,7 @@
 import copy
 import json
 import re
+import subprocess
 import sys
 import warnings
 from pathlib import Path
@@ -395,6 +396,11 @@ def main():
         default=128,
         help="Number of scoring steps for auto quantization. Default is 128.",
     )
+    parser.add_argument(
+        "--trt_build",
+        action="store_true",
+        help="Build a TensorRT engine from the exported ONNX model using trtexec.",
+    )
     parser.add_argument(
         "--no_pretrained",
         action="store_true",
@@ -496,6 +502,26 @@ def main():
 
     print(f"Quantized ONNX model is saved to {args.onnx_save_path}")
 
+    if args.trt_build:
+        build_trt_engine(args.onnx_save_path)
+
+
+def build_trt_engine(onnx_path):
+    """Build a TensorRT engine from the exported ONNX model using trtexec."""
+    cmd = [
+        "trtexec",
+        f"--onnx={onnx_path}",
+        "--stronglyTyped",
+        "--builderOptimizationLevel=4",
+    ]
+    print(f"\nBuilding TensorRT engine: {' '.join(cmd)}")
+    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"TensorRT engine build failed for {onnx_path}:\n{result.stdout}\n{result.stderr}"
+        )
+    print("TensorRT engine build succeeded.")
+
 
 if __name__ == "__main__":
     main()

tests/examples/torch_onnx/test_torch_quant_to_onnx.py

@@ -14,9 +14,6 @@
 # limitations under the License.
 
 
-import os
-import subprocess
-
 import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 
@@ -31,40 +28,13 @@
     "resnet50": ("resnet50", None),
 }
 
-# Builder optimization level: 4 for low-bit modes, 3 otherwise
-_LOW_BIT_MODES = {"fp8", "int8", "nvfp4"}
-
-
-def _verify_trt_engine_build(onnx_save_path, quantize_mode):
-    """Verify the exported ONNX model can be compiled into a TensorRT engine."""
-    example_dir = os.path.join(
-        os.path.dirname(__file__), "..", "..", "..", "examples", "torch_onnx"
-    )
-    onnx_path = os.path.join(example_dir, onnx_save_path)
-    assert os.path.exists(onnx_path), f"ONNX file not found: {onnx_path}"
-
-    opt_level = "4" if quantize_mode in _LOW_BIT_MODES else "3"
-    cmd = [
-        "trtexec",
-        f"--onnx={onnx_path}",
-        "--stronglyTyped",
-        f"--builderOptimizationLevel={opt_level}",
-    ]
-
-    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
-    assert result.returncode == 0, (
-        f"TensorRT engine build failed for {onnx_save_path} "
-        f"(mode={quantize_mode}):\n{result.stdout}\n{result.stderr}"
-    )
-
 
 @pytest.mark.parametrize("quantize_mode", _QUANT_MODES)
 @pytest.mark.parametrize("model_key", list(_MODELS))
 def test_torch_onnx(model_key, quantize_mode):
     timm_model_name, model_kwargs = _MODELS[model_key]
     onnx_save_path = f"{model_key}.{quantize_mode}.onnx"
 
-    # Step 1: Quantize and export to ONNX
     cmd_parts = extend_cmd_parts(
         ["python", "torch_quant_to_onnx.py"],
         timm_model_name=timm_model_name,
@@ -74,8 +44,5 @@ def test_torch_onnx(model_key, quantize_mode):
         calibration_data_size="1",
         num_score_steps="1",
     )
-    cmd_parts.append("--no_pretrained")
+    cmd_parts.extend(["--no_pretrained", "--trt_build"])
     run_example_command(cmd_parts, "torch_onnx")
-
-    # Step 2: Verify the exported ONNX model builds a TensorRT engine
-    _verify_trt_engine_build(onnx_save_path, quantize_mode)