Selaa lähdekoodia

修改剪枝文件解决IR v10版本问题

zhy 2 viikkoa sitten
vanhempi
commit
50f28a858e
1 muutettua tiedostoa jossa 1 lisäystä ja 1 poistoa
  1. 1 1
      tests/prune_tool.py

+ 1 - 1
tests/prune_tool.py

@@ -43,7 +43,7 @@ def prune_weights(model_path, pruned_model, pruning_percentage=0.05):
     )
 
     # 创建新的模型
-    new_model = onnx.helper.make_model(new_graph, producer_name='onnx-example')
+    new_model = onnx.helper.make_model(new_graph, producer_name='onnx-example', opset_imports=model.opset_import, ir_version=model.ir_version)
     # 保存剪枝后的模型
     onnx.save(new_model, pruned_model)