Merge pull request #1 from PINTO0309/fix_irversion

Fix to preserve domain and ir_version
PINTO0309 · Apr 30, 2024 · df14346 · df14346
2 parents 5bab0d3 + 5198075
commit df14346
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/sor4onnx/__init__.py b/sor4onnx/__init__.py
@@ -1,3 +1,3 @@
 from sor4onnx.onnx_opname_renamer import rename, main
 
-__version__ = '1.0.5'
+__version__ = '1.0.6'
diff --git a/sor4onnx/onnx_opname_renamer.py b/sor4onnx/onnx_opname_renamer.py
@@ -136,6 +136,11 @@ def rename(
     # onnx_graph If specified, onnx_graph is processed first
     if not onnx_graph:
         onnx_graph = onnx.load(input_onnx_file_path)
+
+    # domain, ir_version
+    domain: str = onnx_graph.domain
+    ir_version: int = onnx_graph.ir_version
+
     graph = gs.import_onnx(onnx_graph)
     graph.cleanup().toposort()
 
@@ -214,9 +219,9 @@ def rename(
     # Shape Estimation
     renamed_graph = None
     try:
-        renamed_graph = onnx.shape_inference.infer_shapes(gs.export_onnx(graph))
+        renamed_graph = onnx.shape_inference.infer_shapes(gs.export_onnx(graph, do_type_check=False, **{'domain': domain, 'ir_version': ir_version}))
     except:
-        renamed_graph = gs.export_onnx(graph)
+        renamed_graph = gs.export_onnx(graph, do_type_check=False, **{'domain': domain, 'ir_version': ir_version})
         if not non_verbose:
             print(
                 f'{Color.YELLOW}WARNING:{Color.RESET} '+