Merge pull request #271 from YaoHan404/master

fix evaluator/custom.py
huawei-noah · Sep 28, 2022 · 39741b5 · 39741b5
2 parents c1686fb + e511107
commit 39741b5
Showing 1 changed file with 6 additions and 1 deletion.
diff --git 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')