Unverified Commit 37c9d789 authored by SkyTNT's avatar SkyTNT Committed by GitHub
Browse files

Fix is_onnx_available (#440)



* Fix is_onnx_available

Fix: If user install onnxruntime-gpu, is_onnx_available() will return False.

* add more onnxruntime candidates

* Run `make style`
Co-authored-by: default avataranton-l <anton@huggingface.co>
parent 214520c6
...@@ -137,11 +137,19 @@ except importlib_metadata.PackageNotFoundError: ...@@ -137,11 +137,19 @@ except importlib_metadata.PackageNotFoundError:
_onnx_available = importlib.util.find_spec("onnxruntime") is not None _onnx_available = importlib.util.find_spec("onnxruntime") is not None
try: if _onnx_available:
_onnxruntime_version = importlib_metadata.version("onnxruntime") candidates = ("onnxruntime", "onnxruntime-gpu", "onnxruntime-directml", "onnxruntime-openvino")
logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}") _onnxruntime_version = None
except importlib_metadata.PackageNotFoundError: # For the metadata, we have to look for both onnxruntime and onnxruntime-gpu
_onnx_available = False for pkg in candidates:
try:
_onnxruntime_version = importlib_metadata.version(pkg)
break
except importlib_metadata.PackageNotFoundError:
pass
_onnx_available = _onnxruntime_version is not None
if _onnx_available:
logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}")
_scipy_available = importlib.util.find_spec("scipy") is not None _scipy_available = importlib.util.find_spec("scipy") is not None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment