Skip to content

Commit 39741b5

Browse files
authored
Merge pull request #271 from YaoHan404/master
fix evaluator/custom.py
2 parents c1686fb + e511107 commit 39741b5

1 file changed

Lines changed: 6 additions & 1 deletion

File tree

vega/evaluator/custom.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,12 @@ def export_model(self, init_model):
3737
from mindspore.train.serialization import export
3838
from mindspore import Tensor
3939
from mindspore.common.api import _cell_graph_executor
40-
_cell_graph_executor.set_jit_config(jit_config={"jit_level": "o0"})
40+
if hasattr(_cell_graph_executor, "set_jit_config"):
41+
_cell_graph_executor.set_jit_config(jit_config={"jit_level": "o0"})
42+
if hasattr(init_model, "set_jit_config"):
43+
from mindspore.common.jit_conig import JitConfig
44+
jit_conig = JitConfig(jit_level="O0")
45+
init_model.set_jit_config(jit_conig)
4146
fake_input = np.random.random([1, 12, 320, 320]).astype(np.float32)
4247
save_name = os.path.join("./", "ms2air.air")
4348
export(init_model, Tensor(fake_input), Tensor(640), file_name=save_name, file_format='AIR')

0 commit comments

Comments
 (0)