[pytorch] Add macOS M1 chip support. (#1656)

Change-Id: I5fb7f2bc1ba07626fb8dcc45e1daf77d4b31765b
deepjavalibrary · May 16, 2022 · e9a14c8 · e9a14c8
1 parent a8aa45e
commit e9a14c8
Show file tree

Hide file tree

Showing 2 changed files with 3 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
@@ -364,7 +364,8 @@ private static LibTorch downloadPyTorch(Platform platform) {
         String precxx11;
         if (Boolean.getBoolean("PYTORCH_PRECXX11")
                 || Boolean.parseBoolean(System.getenv("PYTORCH_PRECXX11"))
-                || "aarch64".equals(platform.getOsArch())) {
+                || ("aarch64".equals(platform.getOsArch())
+                        && "linux".equals(platform.getOsPrefix()))) {
             precxx11 = "-precxx11";
         } else {
             precxx11 = "";

diff --git a/engines/pytorch/pytorch-jni/build.gradle b/engines/pytorch/pytorch-jni/build.gradle
@@ -30,6 +30,7 @@ processResources {
             files.add("linux-x86_64/cu113/libdjl_torch.so")
             files.add("linux-x86_64/cu113-precxx11/libdjl_torch.so")
             files.add("win-x86_64/cu113/djl_torch.dll")
+            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
         } else if (ptVersion.startsWith("1.10.")) {
             files.add("linux-x86_64/cu113/libdjl_torch.so")
             files.add("linux-x86_64/cu113-precxx11/libdjl_torch.so")