From 7a0f9b60d0ae22eec9d0027c2cfa5524616e37de Mon Sep 17 00:00:00 2001 From: lijiaqi5 Date: Wed, 10 Sep 2025 13:45:43 +0800 Subject: [PATCH] Support provide ORT session options Signed-off-by: opluss --- tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py b/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py index 7884c85b..53f01707 100644 --- a/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py +++ b/tools/onnx-graphsurgeon/onnx_graphsurgeon/ir/graph.py @@ -699,6 +699,7 @@ def fold_constants( size_threshold=None, should_exclude_node=None, recurse_functions=True, + ort_session_options=None, ): """ Folds constants in-place in the graph. The graph's nodes and functions must be topologically @@ -754,7 +755,8 @@ def fold_constants( recurse_functions (bool): Whether to fold constants in this graph's Functions. Defaults to True. - + ort_session_options (Optional[onnxruntime.SessionOptions]): + SessionOptions object to be used for ONNX Runtime sessions. Returns: self """ @@ -1172,6 +1174,7 @@ def get_out_node_ids(): sess = onnxrt.InferenceSession( export_onnx(part, do_type_check=False).SerializeToString(), + sess_options = ort_session_options, providers=ORT_PROVIDERS, ) values = sess.run(names, {}) @@ -1254,6 +1257,7 @@ def should_eval_foldable(tensor): export_onnx( graph_clone, do_type_check=False ).SerializeToString(), + sess_options = ort_session_options, providers=ORT_PROVIDERS, ) values = sess.run(names, {})