[pytorch] Allows to exclude certain DLL from pytorch directory (#3043)

deepjavalibrary · Mar 26, 2024 · 091a41c · 091a41c
1 parent fa28fbe
commit 091a41c
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
@@ -118,6 +118,8 @@ private static void loadLibTorch(LibTorch libTorch) {
                 }
             }
         }
+        String libExclusion = Utils.getEnvOrSystemProperty("PYTORCH_LIBRARY_EXCLUSION", "");
+        Set<String> exclusion = new HashSet<>(Arrays.asList(libExclusion.split(",")));
         boolean isCuda = libTorch.flavor.contains("cu");
         List<String> deferred =
                 Arrays.asList(
@@ -137,7 +139,8 @@ private static void loadLibTorch(LibTorch libTorch) {
             paths.filter(
                             path -> {
                                 String name = path.getFileName().toString();
-                                if (!LIB_PATTERN.matcher(name).matches()) {
+                                if (!LIB_PATTERN.matcher(name).matches()
+                                        || exclusion.contains(name)) {
                                     return false;
                                 } else if (!isCuda
                                         && name.contains("nvrtc")