diff options
| author | Nikita Shulga <nshulga@fb.com> | 2021-12-08 14:59:16 -0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2021-12-08 14:59:16 -0800 |
| commit | c78ceadbb0246ef96743c77b85b21910ca97eb6e (patch) | |
| tree | 8cd2722760f6936709ee305413b0bc2088d31f92 | |
| parent | 70af72c7941995ffe24cfa37f8c8632e5c6135fa (diff) | |
| download | pytorch-c78ceadbb0246ef96743c77b85b21910ca97eb6e.tar.gz pytorch-c78ceadbb0246ef96743c77b85b21910ca97eb6e.tar.bz2 pytorch-c78ceadbb0246ef96743c77b85b21910ca97eb6e.zip | |
[LiteInterpreter] Specify `Loader` to `yaml.load` (#67694) (#69642)
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
| -rw-r--r-- | tools/lite_interpreter/gen_selected_mobile_ops_header.py | 9 |
1 files changed, 8 insertions, 1 deletions
diff --git a/tools/lite_interpreter/gen_selected_mobile_ops_header.py b/tools/lite_interpreter/gen_selected_mobile_ops_header.py index bf28bf3c3a..4b9ee56e58 100644 --- a/tools/lite_interpreter/gen_selected_mobile_ops_header.py +++ b/tools/lite_interpreter/gen_selected_mobile_ops_header.py @@ -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; }""" @@ -121,7 +128,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) |