set pt dir to LD_LIBRARY_PATH

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Jun 21, 2024 · ced2562 · ced2562
1 parent 2e91ce8
commit ced2562
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/deepmd/tf/lmp.py b/deepmd/tf/lmp.py
@@ -6,6 +6,9 @@
 from importlib import (
     import_module,
 )
+from importlib.util import (
+    find_spec,
+)
 from pathlib import (
     Path,
 )
@@ -77,6 +80,11 @@ def get_library_path(module: str, filename: str) -> List[str]:
 
 tf_dir = tf.sysconfig.get_lib()
 op_dir = str(SHARED_LIB_DIR)
+pt_spec = find_spec("torch")
+if pt_spec is not None:
+    pt_dir = pt_spec.submodule_search_locations[0]
+else:
+    pt_dir = None
 
 
 cuda_library_paths = []
@@ -106,6 +114,7 @@ def get_library_path(module: str, filename: str) -> List[str]:
     [
         os.environ.get(lib_env),
         tf_dir,
+        pt_dir,
         os.path.join(tf_dir, "python"),
         op_dir,
     ]