try to add rpath nvidia python

OpenNMT · Oct 29, 2024 · 4ed37d0 · 4ed37d0
1 parent 383d063
commit 4ed37d0
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/python/setup.py b/python/setup.py
@@ -1,6 +1,7 @@
 import glob
 import os
 import sys
+import site
 
 import pybind11
 
@@ -51,6 +52,14 @@ def _maybe_add_library_root(lib_name):
 elif sys.platform == "win32":
     cflags = ["/std:c++17", "/d2FH4-"]
     package_data["ctranslate2"] = ["*.dll"]
+else:
+    CUDA_RPATHS=[
+        '$ORIGIN/../nvidia/cublas/lib',
+        '$ORIGIN/../nvidia/cuda_runtime/lib',
+        '$ORIGIN/../nvidia/cudnn/lib',
+    ]
+    CUDA_RPATHS= ';'.join(map(str, CUDA_RPATHS))
+    ldflags.append("-Wl,-rpath," + CUDA_RPATHS + "/lib")
 
 ctranslate2_module = Extension(
     "ctranslate2._ext",