Fix dnnl wheel package name (#2823)

* Append '-dnnl' to whl package name when --use_dnnl

* Update build.py
This commit is contained in:
Yingge WAN 2020-01-14 06:37:11 +08:00 committed by George Wu
parent 7c6242b024
commit 07502ec14e
2 changed files with 8 additions and 2 deletions

View file

@ -34,6 +34,10 @@ elif '--use_cuda' in sys.argv:
elif '--use_ngraph' in sys.argv:
package_name = 'onnxruntime-ngraph'
sys.argv.remove('--use_ngraph')
elif '--use_dnnl' in sys.argv:
package_name = 'onnxruntime-dnnl'
sys.argv.remove('--use_dnnl')
elif '--use_openvino' in sys.argv:
package_name = 'onnxruntime-openvino'

View file

@ -705,7 +705,7 @@ def nuphar_run_python_tests(build_dir, configs):
run_subprocess([sys.executable, 'onnxruntime_test_python_nuphar.py'], cwd=cwd, dll_path=dll_path)
def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_ngraph, use_tensorrt, use_openvino, use_nuphar, nightly_build = False):
def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_ngraph, use_dnnl, use_tensorrt, use_openvino, use_nuphar, nightly_build = False):
for config in configs:
cwd = get_config_build_dir(build_dir, config)
if is_windows():
@ -719,6 +719,8 @@ def build_python_wheel(source_dir, build_dir, configs, use_cuda, use_ngraph, use
args.append('--use_cuda')
elif use_ngraph:
args.append('--use_ngraph')
elif use_dnnl:
args.append('--use_dnnl')
elif use_openvino:
args.append('--use_openvino')
elif use_nuphar:
@ -969,7 +971,7 @@ def main():
if args.build:
if args.build_wheel:
nightly_build = bool(os.getenv('NIGHTLY_BUILD') == '1')
build_python_wheel(source_dir, build_dir, configs, args.use_cuda, args.use_ngraph, args.use_tensorrt, args.use_openvino, args.use_nuphar, nightly_build)
build_python_wheel(source_dir, build_dir, configs, args.use_cuda, args.use_ngraph, args.use_dnnl, args.use_tensorrt, args.use_openvino, args.use_nuphar, nightly_build)
if args.gen_doc and (args.build or args.test):
generate_documentation(source_dir, build_dir, configs)