Fix to preserve domain and ir_version

PINTO0309 · Apr 30, 2024 · fe71149 · fe71149
1 parent 8085b0c
commit fe71149
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/sam4onnx/__init__.py b/sam4onnx/__init__.py
@@ -1,3 +1,3 @@
 from sam4onnx.onnx_attr_const_modify import modify, main
 
-__version__ = '1.0.14'
+__version__ = '1.0.15'
diff --git a/sam4onnx/onnx_attr_const_modify.py b/sam4onnx/onnx_attr_const_modify.py
@@ -301,6 +301,10 @@ def modify(
                     if node.domain not in ONNX_STANDARD_DOMAINS
             ]
 
+    # domain, ir_version
+    domain: str = onnx_graph.domain
+    ir_version: int = onnx_graph.ir_version
+
     graph = gs.import_onnx(onnx_graph)
 
     # Check if Graph contains a custom domain (custom module)
@@ -436,7 +440,7 @@ def modify(
 
     # Cleanup
     graph.cleanup().toposort()
-    modified_graph = gs.export_onnx(graph)
+    modified_graph = gs.export_onnx(graph, do_type_check=False, **{'domain': domain, 'ir_version': ir_version})
 
     # Optimize
     new_model = None