fix: Fix PTQ export (#3447)

Author: peri044

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":

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],