From e511107b6def66ab6d61aceecedf2f88419d2180 Mon Sep 17 00:00:00 2001 From: YaoHan404 <1601710747@qq.com> Date: Wed, 28 Sep 2022 16:42:29 +0800 Subject: [PATCH] fix evaluator/custom.py set_jit_config() of Mindspore changed --- vega/evaluator/custom.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vega/evaluator/custom.py b/vega/evaluator/custom.py index b3bbb7b7..2a7aadcd 100644 --- a/vega/evaluator/custom.py +++ b/vega/evaluator/custom.py @@ -37,7 +37,12 @@ def export_model(self, init_model): from mindspore.train.serialization import export from mindspore import Tensor from mindspore.common.api import _cell_graph_executor - _cell_graph_executor.set_jit_config(jit_config={"jit_level": "o0"}) + if hasattr(_cell_graph_executor, "set_jit_config"): + _cell_graph_executor.set_jit_config(jit_config={"jit_level": "o0"}) + if hasattr(init_model, "set_jit_config"): + from mindspore.common.jit_conig import JitConfig + jit_conig = JitConfig(jit_level="O0") + init_model.set_jit_config(jit_conig) fake_input = np.random.random([1, 12, 320, 320]).astype(np.float32) save_name = os.path.join("./", "ms2air.air") export(init_model, Tensor(fake_input), Tensor(640), file_name=save_name, file_format='AIR')