From 4eff7b0d6786908c7a2f30e87c10397ed20f7f31 Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Wed, 19 Mar 2025 10:54:21 -0700 Subject: [PATCH] fix: Fix PTQ export --- examples/dynamo/vgg16_ptq.py | 4 +--- tests/py/dynamo/models/test_models_export.py | 3 +-- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/examples/dynamo/vgg16_ptq.py b/examples/dynamo/vgg16_ptq.py index 4ca19e1fd0..7fa943040e 100644 --- a/examples/dynamo/vgg16_ptq.py +++ b/examples/dynamo/vgg16_ptq.py @@ -233,10 +233,8 @@ def calibrate_loop(model): with export_torch_mode(): # Compile the model with Torch-TensorRT Dynamo backend input_tensor = images.cuda() - # torch.export.export() failed due to RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode() - from torch.export._trace import _export - exp_program = _export(model, (input_tensor,)) + exp_program = torch.export.export(model, (input_tensor,), strict=False) if args.quantize_type == "int8": enabled_precisions = {torch.int8} elif args.quantize_type == "fp8": diff --git a/tests/py/dynamo/models/test_models_export.py b/tests/py/dynamo/models/test_models_export.py index 2d0992ca8b..469ed569d1 100644 --- a/tests/py/dynamo/models/test_models_export.py +++ b/tests/py/dynamo/models/test_models_export.py @@ -206,7 +206,6 @@ def test_resnet18_half(ir): def test_base_fp8(ir): import modelopt.torch.quantization as mtq from modelopt.torch.quantization.utils import export_torch_mode - from torch.export._trace import _export class SimpleNetwork(torch.nn.Module): def __init__(self): @@ -234,7 +233,7 @@ def calibrate_loop(model): with torch.no_grad(): with export_torch_mode(): - exp_program = _export(model, (input_tensor,)) + exp_program = torch.export.export(model, (input_tensor,), strict=False) trt_model = torchtrt.dynamo.compile( exp_program, inputs=[input_tensor],