frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[LiteInterpreter] Specify Loader to yaml.load (#67694)

Browse Source 
Summary:
It became a mandatory argument since PyYaml-6, but has been present since PyYaml-3

Unblock migration to newer runtime

Pull Request resolved: https://github.com/pytorch/pytorch/pull/67694

Reviewed By: seemethere

Differential Revision: D32106043

Pulled By: malfet

fbshipit-source-id: 35246b97a974b168c066396ea31987b267534c7f
This commit is contained in:
Nikita Shulga
2021-11-02 12:51:16 -07:00
committed by Facebook GitHub Bot
parent b00206d473
commit 8f63cfda14
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tools/lite_interpreter/gen_selected_mobile_ops_header.py
View File

@ -7,6 +7,13 @@ from tools.codegen.code_template import CodeTemplate
import yaml
# Safely load fast C Yaml loader/dumper if they are available
try:
from yaml import CSafeLoader as Loader
except ImportError:
from yaml import SafeLoader as Loader # type: ignore[misc]
if_condition_template_str = """if (kernel_tag_sv.compare("$kernel_tag_name") == 0) {
return $dtype_checks;
}"""
@ -127,7 +134,7 @@ def main() -> None:
print("Loading yaml file: ", model_file_name)
loaded_model = {}
with open(model_file_name, "rb") as model_file:
loaded_model = yaml.load(model_file)
loaded_model = yaml.load(model_file, Loader=Loader)
root_operators_set = set(loaded_model)