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 exclusion = new HashSet<>(Arrays.asList(libExclusion.split(","))); boolean isCuda = libTorch.flavor.contains("cu"); List 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")