From 091a41cbba49dc8dc7d90f52014d40e13011fb99 Mon Sep 17 00:00:00 2001
From: Frank Liu <frankfliu2000@gmail.com>
Date: Tue, 26 Mar 2024 10:08:26 -0700
Subject: [PATCH] [pytorch] Allows to exclude certain DLL from pytorch
 directory (#3043)

---
 .../src/main/java/ai/djl/pytorch/jni/LibUtils.java           | 5 ++++-
 1 file changed, 4 insertions(+), 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
index 8dcd5c7aa72..83e0f5b5b95 100644
--- 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")