[pytorch] Avoid loading .lib file from PYTORCH_LIBRARY_PATH (#2987)

deepjavalibrary · Feb 11, 2024 · cf46182 · cf46182
1 parent f9ad5f1
commit cf46182
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
@@ -65,6 +65,7 @@ public final class LibUtils {
 
  private static final Pattern VERSION_PATTERN =
  Pattern.compile("(\\d+\\.\\d+\\.\\d+(-[a-z]+)?)(-SNAPSHOT)?(-\\d+)?");
+ private static final Pattern LIB_PATTERN = Pattern.compile("(.*\\.(so(\\.\\d+)*|dll|dylib))");
 
  private static LibTorch libTorch;
 
@@ -136,7 +137,9 @@ private static void loadLibTorch(LibTorch libTorch) {
  paths.filter(
  path -> {
  String name = path.getFileName().toString();
- if (!isCuda
+ if (!LIB_PATTERN.matcher(name).matches()) {
+ return false;
+ } else if (!isCuda
  && name.contains("nvrtc")
  && name.contains("cudart")
  && name.contains("nvTools")) {