From 2a8cec57027cf7d1c9eaf5cc7a6e5d33adaa88c8 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Tue, 12 Mar 2024 13:55:21 -0700 Subject: [PATCH] [pytorch] Fixes detecting wrong flavor on macOS issue (#3027) --- .../main/java/ai/djl/pytorch/jni/LibUtils.java | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) 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 4d2e8505f02..8dcd5c7aa72 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 @@ -231,10 +231,21 @@ private static Path findJniLibrary(LibTorch libTorch) { String djlVersion = libTorch.apiVersion; String flavor = libTorch.flavor; + // Looking for JNI in libTorch.dir first + Path libDir = libTorch.dir.toAbsolutePath(); + Path path = libDir.resolve(djlVersion + '-' + JNI_LIB_NAME); + if (Files.exists(path)) { + return path; + } + path = libDir.resolve(JNI_LIB_NAME); + if (Files.exists(path)) { + return path; + } + // always use cache dir, cache dir might be different from libTorch.dir Path cacheDir = Utils.getEngineCacheDir("pytorch"); Path dir = cacheDir.resolve(version + '-' + flavor + '-' + classifier); - Path path = dir.resolve(djlVersion + '-' + JNI_LIB_NAME); + path = dir.resolve(djlVersion + '-' + JNI_LIB_NAME); if (Files.exists(path)) { return path; } @@ -554,8 +565,10 @@ private static final class LibTorch { if (flavor == null || flavor.isEmpty()) { if (CudaUtils.getGpuCount() > 0) { flavor = "cu" + CudaUtils.getCudaVersionString() + "-precxx11"; - } else { + } else if ("linux".equals(platform.getOsPrefix())) { flavor = "cpu-precxx11"; + } else { + flavor = "cpu"; } } }