Update export function for torch 2.6 (#2184)

Signed-off-by: Yi Liu <yiliu4@habana.ai>

Author: yiliu30

neural_compressor/torch/export/pt2e_export.py

@@ -16,11 +16,10 @@
 from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
-from torch._export import capture_pre_autograd_graph
 from torch.fx.graph_module import GraphModule
 
 from neural_compressor.common.utils import logger
-from neural_compressor.torch.utils import TORCH_VERSION_2_2_2, get_torch_version, is_ipex_imported
+from neural_compressor.torch.utils import TORCH_VERSION_2_2_2, TORCH_VERSION_2_7_0, get_torch_version, is_ipex_imported
 
 __all__ = ["export", "export_model_for_pt2e_quant"]
 
@@ -53,6 +52,10 @@ def export_model_for_pt2e_quant(
             # Note 1: `capture_pre_autograd_graph` is also a short-term API, it will be
             # updated to use the official `torch.export` API when that is ready.
             cur_version = get_torch_version()
+            if cur_version >= TORCH_VERSION_2_7_0:
+                export_func = torch.export.export_for_training
+            else:
+                export_func = torch._export.capture_pre_autograd_graph
             if cur_version <= TORCH_VERSION_2_2_2:  # pragma: no cover
                 logger.warning(
                     (
@@ -62,11 +65,13 @@ def export_model_for_pt2e_quant(
                     ),
                     cur_version,
                 )
-                exported_model = capture_pre_autograd_graph(model, args=example_inputs)
+                exported_model = export_func(model, args=example_inputs)
             else:
-                exported_model = capture_pre_autograd_graph(  # pylint: disable=E1123
+                exported_model = export_func(  # pylint: disable=E1123
                     model, args=example_inputs, dynamic_shapes=dynamic_shapes
                 )
+            if cur_version >= TORCH_VERSION_2_7_0:
+                exported_model = exported_model.module()
             exported_model._exported = True
             logger.info("Exported the model to Aten IR successfully.")
     except Exception as e:

neural_compressor/torch/utils/environ.py

@@ -129,6 +129,7 @@ def get_ipex_version():
 
 
 TORCH_VERSION_2_2_2 = Version("2.2.2")
+TORCH_VERSION_2_7_0 = Version("2.7.0")
 
 
 def get_torch_version():