[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")) {