Skip to content

Commit

Permalink
Fix path used in check for cudnn library (#7786)
Browse files Browse the repository at this point in the history
* There are separate paths for CUDA and CUDNN as they are not guaranteed to be in the same location on a Windows machine. Use the CUDNN path when looking for the CUDNN library.

* Refine check
  • Loading branch information
skottmckay authored May 27, 2021
1 parent ddf4aaa commit 63df683
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions onnxruntime/python/_pybind_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,16 @@
raise ImportError(f"CUDA Toolkit {cuda_version_major}.x not installed on the machine.")

cuda_bin_dir = os.path.join(os.environ[cuda_env_variable], "bin")
if not os.path.isfile(os.path.join(cuda_bin_dir, f"cudnn64_{version_info.cudnn_version}.dll")):
raise ImportError(f"cuDNN {version_info.cudnn_version} not installed in {cuda_bin_dir}.")

# prefer CUDNN_HOME if set. fallback to the CUDA install directory (would have required user to manually
# copy the cudnn dll there
cudnn_path = os.environ["CUDNN_HOME"] if "CUDNN_HOME" in os.environ else os.environ[cuda_env_variable]
cudnn_bin_dir = os.path.join(cudnn_path, "bin")

if not os.path.isfile(os.path.join(cudnn_bin_dir, f"cudnn64_{version_info.cudnn_version}.dll")):
raise ImportError(f"cuDNN {version_info.cudnn_version} not installed in {cudnn_bin_dir}. "
f"Set the CUDNN_HOME environment variable to the path of the 'cuda' directory "
f"in your CUDNN installation if necessary.")

if sys.version_info >= (3, 8):
# Python 3.8 (and later) doesn't search system PATH when loading DLLs, so the CUDA location needs to be
Expand Down

0 comments on commit 63df683

Please sign in to comment.