nit: added ir_version as arg in model building function

Signed-off-by: gcunhase <[email protected]>

Author: gcunhase

tests/_test_utils/onnx_quantization/lib_test_models.py

@@ -556,7 +556,7 @@ def build_convtranspose_conv_residual_model():
     return model_inferred
 
 
-def build_matmul_relu_model_ir_12():
+def build_matmul_relu_model(ir_version=12):
     # Define your model inputs and outputs
     input_names = ["input_0"]
     output_names = ["output_0"]
@@ -599,12 +599,14 @@ def build_matmul_relu_model_ir_12():
     ]
 
     # Create the ONNX graph with the nodes and initializers
-    graph = helper.make_graph(nodes, "matmul_relu", inputs, outputs, initializer=initializers)
+    graph = helper.make_graph(
+        nodes, f"matmul_relu_ir_{ir_version}", inputs, outputs, initializer=initializers
+    )
 
     # Create the ONNX model
     model = helper.make_model(graph)
     model.opset_import[0].version = 13
-    model.ir_version = 12
+    model.ir_version = ir_version
 
     # Check the ONNX model
     model_inferred = onnx.shape_inference.infer_shapes(model)

tests/unit/onnx/test_onnx_utils.py

@@ -18,7 +18,7 @@
 import numpy as np
 import onnx
 import pytest
-from _test_utils.onnx_quantization.lib_test_models import build_matmul_relu_model_ir_12
+from _test_utils.onnx_quantization.lib_test_models import build_matmul_relu_model
 from _test_utils.torch_model.vision_models import get_tiny_resnet_and_input
 from onnx.helper import (
     make_graph,
@@ -258,7 +258,7 @@ def test_remove_node_extra_training_outputs():
 
 
 def test_ir_version_support(tmp_path):
-    model = build_matmul_relu_model_ir_12()
+    model = build_matmul_relu_model(ir_version=12)
     model_path = os.path.join(tmp_path, "test_matmul_relu.onnx")
     onnx.save(model, model_path)
     model_reload, _, _, _, _ = load_onnx_model(model_path, intermediate_generated_files=[])