Repository: deepjavalibrary/djl
Pull request #3163 (2024-05-07 18:34:28)

[sentencepiece] Updates sentencepiece to 0.2.0
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index dfb33432b..b31bf97ea 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -52,11 +52,11 @@ jobs:
   build-pytorch-jni-linux:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: ubuntu-latest
-    container: nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04
+    container: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
     steps:
       - uses: actions/checkout@v3
       - name: Set up JDK 17
-        uses: actions/setup-java@v3
+        uses: actions/setup-java@v4
         with:
           distribution: 'corretto'
           java-version: 17
@@ -70,11 +70,7 @@ jobs:
         run: |
           apt-get update
           DEBIAN_FRONTEND=noninteractive apt-get install -y locales curl unzip software-properties-common
-          add-apt-repository -y ppa:deadsnakes/ppa
-          apt-get update
-          apt-get install -y python3 python3-distutils
-          curl -O https://bootstrap.pypa.io/pip/3.6/get-pip.py
-          python3 get-pip.py
+          apt-get install -y python3-pip python3-distutils
           pip3 install awscli cmake
           ln -s /usr/local/bin/cmake /usr/bin/cmake
       - name: Release JNI prep
@@ -86,10 +82,9 @@ jobs:
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=$PYTORCH_VERSION
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
-          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Ppt_version=$PYTORCH_VERSION
+          export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
+          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=cu121 -Ppt_version=$PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
-          if [[ "$PYTORCH_VERSION" == "1.12.1" ]]; then ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu10 -Ppt_version=$PYTORCH_VERSION; fi
-          if [[ "$PYTORCH_VERSION" == "1.11.0" ]]; then ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu10 -Ppt_version=$PYTORCH_VERSION; fi
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v3
         with:
@@ -106,13 +101,13 @@ jobs:
   build-pytorch-jni-precxx11:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: ubuntu-latest
-    container: nvidia/cuda:11.3.1-cudnn8-devel-centos7
+    container: nvidia/cuda:12.1.1-cudnn8-devel-centos7
     steps:
       - name: Install Environment
         run: |
           yum -y update
           yum -y install centos-release-scl-rh epel-release
-          yum -y install devtoolset-7 rh-git218 patch cmake3
+          yum -y install devtoolset-9 rh-git218 patch cmake3
           ln -s /usr/bin/cmake3 /usr/bin/cmake
           pip3 install awscli --upgrade
       - uses: actions/checkout@v3
@@ -129,7 +124,7 @@ jobs:
             ${{ runner.os }}-gradle-
       - name: Release JNI prep
         run: |
-          source scl_source enable devtoolset-7 rh-git218 || true
+          source scl_source enable devtoolset-9 rh-git218 || true
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
           export PYTORCH_VERSION=${PYTORCH_VERSION:-$(cat gradle.properties | awk -F '=' '/pytorch_version/ {print $2}')}
           export PYTORCH_PRECXX11=true
@@ -138,7 +133,8 @@ jobs:
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           rm -rf ~/.djl.ai
-          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
+          export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
+          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=cu121 -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v2
         with:
@@ -189,7 +185,7 @@ jobs:
           call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
           set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v11.7"
           set "PATH=%CUDA_PATH%\bin;%CUDA_PATH%\libnvvp;%PATH%"
-          gradlew :engines:pytorch:pytorch-native:cleanJNI :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Ppt_version=${{ github.event.inputs.pt_version }}
+          gradlew :engines:pytorch:pytorch-native:cleanJNI :engines:pytorch:pytorch-native:compileJNI -Pcuda=cu121 -Ppt_version=${{ github.event.inputs.pt_version }}
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v4
         with:
@@ -269,8 +265,13 @@ jobs:
       - name: Install compiler environment
         run: |
           yum -y update
-          yum -y groupinstall "Development Tools"
           yum -y install patch git cmake3 python3-devel java-17-amazon-corretto-devel
+          yum -y install gcc10 gcc10-c++
+          ln -sf /usr/bin/gcc10-gcc /usr/bin/gcc
+          ln -sf /usr/bin/gcc10-cc /usr/bin/cc
+          ln -sf /usr/bin/gcc10-g++ /usr/bin/g++
+          ln -sf /usr/bin/gcc10-c++ /usr/bin/c++
+          ln -sf /usr/bin/gcc10-ar /usr/bin/ar
           ln -sf /usr/bin/cmake3 /usr/bin/cmake
           pip3 install awscli --upgrade
       - uses: actions/checkout@v3
diff --git a/.github/workflows/native_s3_sentencepiece.yml b/.github/workflows/native_s3_sentencepiece.yml
index aa1f543a2..84afdf390 100644
--- a/.github/workflows/native_s3_sentencepiece.yml
+++ b/.github/workflows/native_s3_sentencepiece.yml
@@ -54,7 +54,7 @@ jobs:
         with:
           distribution: 'corretto'
           java-version: 17
-      - uses: actions/cache@v4
+      - uses: actions/cache@v3
         with:
           path: ~/.gradle/caches
           key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
@@ -170,7 +170,7 @@ jobs:
           java-version: 17
           distribution: corretto
           architecture: aarch64
-      - uses: actions/cache@v4
+      - uses: actions/cache@v3
         with:
           path: ~/.gradle/caches
           key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
diff --git a/.github/workflows/nightly_publish.yml b/.github/workflows/nightly_publish.yml
index 48e579590..6d303bbf7 100644
--- a/.github/workflows/nightly_publish.yml
+++ b/.github/workflows/nightly_publish.yml
@@ -190,6 +190,7 @@ jobs:
         run: |
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.13.1 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.1.2 -Psnapshot
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.2.2 -Psnapshot
           ./gradlew clean engines:ml:xgboost:publish -Pgpu -Psnapshot
           ./gradlew clean publish -Psnapshot
           cd bom
@@ -204,6 +205,7 @@ jobs:
         run: |
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.13.1 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.1.2 -P${{ github.event.inputs.mode }}
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.2.2 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:ml:xgboost:publish -Pgpu -P${{ github.event.inputs.mode }}
           ./gradlew clean publish -P${{ github.event.inputs.mode }}
           cd bom
diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
index a0d246626..80779067c 100644
--- a/engines/pytorch/pytorch-engine/README.md
+++ b/engines/pytorch/pytorch-engine/README.md
@@ -24,13 +24,13 @@ The javadocs output is built in the `build/doc/javadoc` folder.
 ## Installation
 You can pull the PyTorch engine from the central Maven repository by including the following dependency:
 
-- ai.djl.pytorch:pytorch-engine:0.27.0
+- ai.djl.pytorch:pytorch-engine:0.28.0
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-engine</artifactId>
-    <version>0.27.0</version>
+    <version>0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -46,7 +46,7 @@ The following table illustrates which pytorch version that DJL supports:
 
 | PyTorch engine version | PyTorch native library version            |
 |------------------------|-------------------------------------------|
-| pytorch-engine:0.28.0  | 1.13.1, **2.1.2**                         |
+| pytorch-engine:0.28.0  | 1.13.1, 2.1.2, **2.2.2**                  |
 | pytorch-engine:0.27.0  | 1.13.1, **2.1.1**                         |
 | pytorch-engine:0.26.0  | 1.13.1, 2.0.1, **2.1.1**                  |
 | pytorch-engine:0.25.0  | 1.11.0, 1.12.1, **1.13.1**, 2.0.1         |
@@ -115,7 +115,7 @@ export PYTORCH_FLAVOR=cpu
 ### macOS
 For macOS, you can use the following library:
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
 - ai.djl.pytorch:pytorch-native-cpu:2.1.1:osx-x86_64
 
 ```xml
@@ -123,13 +123,13 @@ For macOS, you can use the following library:
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>osx-x86_64</classifier>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -139,21 +139,21 @@ For macOS, you can use the following library:
 ### macOS M1
 For macOS M1, you can use the following library:
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cpu:2.1.1:osx-aarch64
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cpu:2.2.2:osx-aarch64
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>osx-aarch64</classifier>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -164,29 +164,29 @@ installed on your GPU machine, you can use one of the following library:
 
 #### Linux GPU
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cu121:2.1.1:linux-x86_64 - CUDA 12.1
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cu121:2.2.2:linux-x86_64 - CUDA 12.1
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cu121</artifactId>
     <classifier>linux-x86_64</classifier>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
 
 ### Linux CPU
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cpu:2.1.1:linux-x86_64
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cpu:2.2.2:linux-x86_64
 
 ```xml
 <dependency>
@@ -194,20 +194,20 @@ installed on your GPU machine, you can use one of the following library:
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>linux-x86_64</classifier>
     <scope>runtime</scope>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
 
 ### For aarch64 build
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.1.1:linux-aarch64
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.2.2:linux-aarch64
 
 ```xml
 <dependency>
@@ -215,12 +215,12 @@ installed on your GPU machine, you can use one of the following library:
     <artifactId>pytorch-native-cpu-precxx11</artifactId>
     <classifier>linux-aarch64</classifier>
     <scope>runtime</scope>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -230,22 +230,22 @@ installed on your GPU machine, you can use one of the following library:
 We also provide packages for the system like CentOS 7/Ubuntu 14.04 with GLIBC >= 2.17.
 All the package were built with GCC 7, we provided a newer `libstdc++.so.6.24` in the package that contains `CXXABI_1.3.9` to use the package successfully.
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cu121-precxx11:2.1.1:linux-x86_64 - CUDA 12.1
-- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.1.1:linux-x86_64   - CPU
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cu121-precxx11:2.2.2:linux-x86_64 - CUDA 12.1
+- ai.djl.pytorch:pytorch-native-cpu-precxx11:2.2.2:linux-x86_64   - CPU
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cu121-precxx11</artifactId>
     <classifier>linux-x86_64</classifier>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -255,13 +255,13 @@ All the package were built with GCC 7, we provided a newer `libstdc++.so.6.24` i
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cpu-precxx11</artifactId>
     <classifier>linux-x86_64</classifier>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
@@ -276,29 +276,29 @@ For the Windows platform, you can choose between CPU and GPU.
 
 #### Windows GPU
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cu121:2.1.1:win-x86_64 - CUDA 12.1
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cu121:2.2.2:win-x86_64 - CUDA 12.1
 
 ```xml
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-native-cu121</artifactId>
     <classifier>win-x86_64</classifier>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
     <scope>runtime</scope>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
 
 ### Windows CPU
 
-- ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0
-- ai.djl.pytorch:pytorch-native-cpu:2.1.1:win-x86_64
+- ai.djl.pytorch:pytorch-jni:2.2.2-0.28.0
+- ai.djl.pytorch:pytorch-native-cpu:2.2.2:win-x86_64
 
 ```xml
 <dependency>
@@ -306,12 +306,12 @@ For the Windows platform, you can choose between CPU and GPU.
     <artifactId>pytorch-native-cpu</artifactId>
     <classifier>win-x86_64</classifier>
     <scope>runtime</scope>
-    <version>2.1.1</version>
+    <version>2.2.2</version>
 </dependency>
 <dependency>
     <groupId>ai.djl.pytorch</groupId>
     <artifactId>pytorch-jni</artifactId>
-    <version>2.1.1-0.27.0</version>
+    <version>2.2.2-0.28.0</version>
     <scope>runtime</scope>
 </dependency>
 ```
diff --git a/engines/pytorch/pytorch-jni/build.gradle b/engines/pytorch/pytorch-jni/build.gradle
index c2b0ee9dc..4a7865288 100644
--- a/engines/pytorch/pytorch-jni/build.gradle
+++ b/engines/pytorch/pytorch-jni/build.gradle
@@ -21,39 +21,25 @@ processResources {
         def files = [
                 "linux-x86_64/cpu/libdjl_torch.so",
                 "linux-x86_64/cpu-precxx11/libdjl_torch.so",
-                "osx-x86_64/cpu/libdjl_torch.dylib",
+                "linux-aarch64/cpu-precxx11/libdjl_torch.so",
+                "osx-aarch64/cpu/libdjl_torch.dylib",
                 "win-x86_64/cpu/djl_torch.dll"
         ]
-        if (ptVersion.startsWith("2.1.")) {
-            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
+        if (ptVersion.startsWith("2.1.") || ptVersion.startsWith("2.2.")) {
             files.add("linux-x86_64/cu121/libdjl_torch.so")
             files.add("linux-x86_64/cu121-precxx11/libdjl_torch.so")
             files.add("win-x86_64/cu121/djl_torch.dll")
-            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
+            files.add("osx-x86_64/cpu/libdjl_torch.dylib")
         } else if (ptVersion.startsWith("2.0.")) {
-            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
             files.add("linux-x86_64/cu118/libdjl_torch.so")
             files.add("linux-x86_64/cu118-precxx11/libdjl_torch.so")
             files.add("win-x86_64/cu118/djl_torch.dll")
-            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
+            files.add("osx-x86_64/cpu/libdjl_torch.dylib")
         } else if (ptVersion.startsWith("1.13.")) {
-            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
             files.add("linux-x86_64/cu117/libdjl_torch.so")
             files.add("linux-x86_64/cu117-precxx11/libdjl_torch.so")
             files.add("win-x86_64/cu117/djl_torch.dll")
-            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
-        } else if (ptVersion.startsWith("1.12.")) {
-            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
-            files.add("linux-x86_64/cu116/libdjl_torch.so")
-            files.add("linux-x86_64/cu116-precxx11/libdjl_torch.so")
-            files.add("win-x86_64/cu116/djl_torch.dll")
-            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
-        } else if (ptVersion.startsWith("1.11.")) {
-            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
-            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")
+            files.add("osx-x86_64/cpu/libdjl_torch.dylib")
         } else {
             throw new GradleException("Unsupported version: ${ptVersion}.")
         }
diff --git a/engines/pytorch/pytorch-native/build.gradle b/engines/pytorch/pytorch-native/build.gradle
index 99a658bf3..99e9a01a0 100644
--- a/engines/pytorch/pytorch-native/build.gradle
+++ b/engines/pytorch/pytorch-native/build.gradle
@@ -15,22 +15,8 @@ boolean isPrecxx11 = project.hasProperty("precxx11")
 boolean isAarch64 = project.hasProperty("aarch64") || System.properties["os.arch"] == "aarch64"
 
 String FLAVOR = "cpu"
-if (project.hasProperty("cu11")) {
-    if (VERSION.startsWith("1.11.")) {
-        FLAVOR = "cu113"
-    } else if (VERSION.startsWith("1.12.")) {
-        FLAVOR = "cu116"
-    } else if (VERSION.startsWith("1.13.")) {
-        FLAVOR = "cu117"
-    } else if (VERSION.startsWith("2.0.")) {
-        FLAVOR = "cu118"
-    } else if (VERSION.startsWith("2.1.")) {
-        FLAVOR = "cu121"
-    } else {
-        throw new GradleException("Unsupported PyTorch version: ${VERSION}")
-    }
-} else if (project.hasProperty("cu10")) {
-    FLAVOR = "cu102"
+if (project.hasProperty("cuda")) {
+    FLAVOR = project.property("cuda")
 }
 String BINARY_ROOT = "${project.buildDir}/download"
 
@@ -90,24 +76,12 @@ def prepareNativeLib(String binaryRoot, String ver) {
 
     def officialPytorchUrl = "https://download.pytorch.org/libtorch"
     def aarch64PytorchUrl = "https://djl-ai.s3.amazonaws.com/publish/pytorch"
-    String cuda
-    if (ver.startsWith("1.11.")) {
-        cuda = "cu113"
-    } else if (ver.startsWith("1.12.")) {
-        cuda = "cu116"
-    } else if (ver.startsWith("1.13.")) {
-        cuda = "cu117"
-    } else if (ver.startsWith("2.0.")) {
-        cuda = "cu118"
-    } else if (ver.startsWith("2.1.")) {
-        cuda = "cu121"
-    } else {
-        throw new GradleException("Unsupported PyTorch version: ${ver}")
-    }
+    String cuda = "cu121"
 
     def files = [
             "cpu/libtorch-cxx11-abi-shared-with-deps-${ver}%2Bcpu.zip"        : "cpu/linux-x86_64",
-            "cpu/libtorch-macos-${ver}.zip"                                   : "cpu/osx-x86_64",
+            "cpu/libtorch-macos-x86_64-${ver}.zip"                            : "cpu/osx-x86_64",
+            "cpu/libtorch-macos-arm64-${ver}.zip"                             : "cpu/osx-aarch64",
             "cpu/libtorch-win-shared-with-deps-${ver}%2Bcpu.zip"              : "cpu/win-x86_64",
             "${cuda}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cuda}.zip": "${cuda}/linux-x86_64",
             "${cuda}/libtorch-win-shared-with-deps-${ver}%2B${cuda}.zip"      : "${cuda}/win-x86_64",
@@ -117,7 +91,6 @@ def prepareNativeLib(String binaryRoot, String ver) {
 
     def aarch64Files = [
             "${ver}/libtorch-shared-with-deps-${ver}-aarch64.zip": "cpu-precxx11/linux-aarch64",
-            "${ver}/libtorch-macos-${ver}-aarch64.zip"           : "cpu/osx-aarch64"
     ]
 
     copyNativeLibToOutputDir(files, binaryRoot, officialPytorchUrl)
@@ -129,6 +102,12 @@ def prepareNativeLib(String binaryRoot, String ver) {
     exec {
         commandLine 'install_name_tool', '-add_rpath', '@loader_path', "${binaryRoot}/cpu/osx-x86_64/native/lib/libtorch.dylib"
     }
+    exec {
+        commandLine 'install_name_tool', '-add_rpath', '@loader_path', "${binaryRoot}/cpu/osx-aarch64/native/lib/libtorch_cpu.dylib"
+    }
+    exec {
+        commandLine 'install_name_tool', '-add_rpath', '@loader_path', "${binaryRoot}/cpu/osx-aarch64/native/lib/libtorch.dylib"
+    }
 }
 
 def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot, String url) {
@@ -147,7 +126,7 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
 
         copy {
             from("${outputDir}/libtorch/lib/") {
-                include "libarm_compute*", "libc10_cuda.so", "libc10.*", "libcaffe2_nvrtc.so", "libcu*", "libgfortran-*", "libgomp*", "libiomp*", "libnv*", "libopenblasp-*", "libtorch_cpu.*", "libtorch_cuda*.so", "libtorch.*", "asmjit.dll", "c10_cuda.dll", "c10.dll", "caffe2_nvrtc.dll", "cu*.dll", "fbgemm.dll", "nv*.dll", "torch_cpu.dll", "torch_cuda*.dll", "torch.dll", "uv.dll", "zlibwapi.dll"
+                include "libarm_compute*", "libc10_cuda.so", "libc10.*", "libcaffe2_nvrtc.so", "libcu*", "libgfortran-*", "libgomp*", "libiomp*", "libomp-*", "libnv*", "libopenblasp-*", "libtorch_cpu.*", "libtorch_cuda*.so", "libtorch.*", "asmjit.dll", "c10_cuda.dll", "c10.dll", "caffe2_nvrtc.dll", "cu*.dll", "fbgemm.dll", "nv*.dll", "torch_cpu.dll", "torch_cuda*.dll", "torch.dll", "uv.dll", "zlibwapi.dll"
             }
             into("${outputDir}/native/lib")
         }
@@ -163,6 +142,14 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
                 i -> libstd.withOutputStream { it << i }
             }
         }
+        if (entry.value.contains("osx-aarch64")) {
+            def libomp = new File("${outputDir}/native/lib/libomp.dylib")
+            def ompUrl = "https://publish.djl.ai/extra/macos-arm64/libomp.dylib"
+            new URL(ompUrl).withInputStream {
+                i -> libomp.withOutputStream { it << i }
+            }
+        }
+
         delete(file)
         delete("${outputDir}/libtorch")
     }
diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
index ae0456bec..214c38a25 100755
--- a/engines/pytorch/pytorch-native/build.sh
+++ b/engines/pytorch/pytorch-native/build.sh
@@ -29,16 +29,23 @@ if [[ ! -d "libtorch" ]]; then
     fi
 
     if [[ $ARCH == 'aarch64' ]]; then
-      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch${AARCH64_CXX11ABI}-shared-with-deps-${VERSION}-aarch64.zip | jar xv > /dev/null
+      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch${AARCH64_CXX11ABI}-shared-with-deps-${VERSION}-aarch64.zip | jar xv >/dev/null
     else
-      curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv > /dev/null
+      curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv >/dev/null
     fi
-
   elif [[ $PLATFORM == 'darwin' ]]; then
-    if [[ $ARCH == 'aarch64' ]]; then
-      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch-macos-${VERSION}-aarch64.zip | jar xv > /dev/null
+    if [[ "$VERSION" =~ ^(2.2.)* ]]; then
+      if [[ $ARCH == 'aarch64' ]]; then
+        curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-${VERSION}.zip | jar xv >/dev/null
+      else
+        curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-x86_64-${VERSION}.zip | jar xv >/dev/null
+      fi
     else
-      curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-${VERSION}.zip | jar xv > /dev/null
+      if [[ $ARCH == 'aarch64' ]]; then
+        curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch-macos-${VERSION}-aarch64.zip | jar xv >/dev/null
+      else
+        curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-${VERSION}.zip | jar xv >/dev/null
+      fi
     fi
   else
     echo "$PLATFORM is not supported."
@@ -46,10 +53,6 @@ if [[ ! -d "libtorch" ]]; then
   fi
 fi
 
-if [[ "$VERSION" == "1.11.0" ]]; then
-  PT_VERSION=V1_11_X
-fi
-
 if [[ "$FLAVOR" = cu* ]]; then
   USE_CUDA=1
 fi
diff --git a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
index 0c75ae1e9..968368180 100644
--- a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
+++ b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
@@ -14,12 +14,11 @@
 #define DJL_TORCH_DJL_PYTORCH_JNI_UTILS_H
 
 #include <c10/util/typeid.h>
-#include <c10/util/variant.h>
 #include <djl/utils.h>
 #include <jni.h>
 #include <torch/csrc/api/include/torch/enum.h>
 #include <torch/script.h>
-
+#include <variant>
 #include <iostream>
 
 #include "djl_pytorch_jni_log.h"
@@ -30,16 +29,10 @@ namespace utils {
 
 #if !defined(__ANDROID__)
 // for image interpolation
-#ifdef V1_10_X
-typedef torch::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
-    torch::enumtype::kBicubic, torch::enumtype::kTrilinear, torch::enumtype::kArea>
-    mode_t;
-#else
-typedef torch::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
+typedef std::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
     torch::enumtype::kBicubic, torch::enumtype::kTrilinear, torch::enumtype::kArea, torch::enumtype::kNearestExact>
     mode_t;
 #endif
-#endif
 
 inline jint GetDTypeFromScalarType(const torch::ScalarType& type) {
   if (torch::kFloat32 == type) {
@@ -118,10 +111,6 @@ inline mode_t GetInterpolationMode(jint jmode) {
       return torch::kTrilinear;
     case 5:
       return torch::kArea;
-    case 6:
-#ifndef V1_10_X
-      return torch::kNearestExact;
-#endif
     default:
       throw;
   }
diff --git a/gradle.properties b/gradle.properties
index af55e87ed..01ae05892 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -13,11 +13,11 @@ systemProp.org.gradle.internal.publish.checksums.insecure=true
 
 djl_version=0.28.0
 mxnet_version=1.9.1
-pytorch_version=2.1.2
+pytorch_version=2.2.2
 tensorflow_version=2.10.1
 trt_version=8.4.1
 onnxruntime_version=1.17.1
-sentencepiece_version=0.1.97
+sentencepiece_version=0.2.0
 tokenizers_version=0.19.1
 llamacpp_version=b1696
 fasttext_version=0.9.2
@@ -44,6 +44,6 @@ spark_version=3.3.2
 openpnp_opencv_version=4.7.0-0
 
 antlr_version=4.11.1
-testng_version=7.9.0
+testng_version=7.10.2
 junit_version=4.13.2
 mockito_version=5.11.0


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java
index f6cfda911..569c83e99 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java
@@ -25,14 +25,12 @@ public class ALibUtilsTest {
     public void setup() {
         System.setProperty("ai.djl.pytorch.native_helper", ALibUtilsTest.class.getName());
         System.setProperty("STDCXX_LIBRARY_PATH", "/usr/lib/non-exists");
-        System.setProperty("PYTORCH_PRECXX11", "true");
     }
 
     @AfterClass
     public void teardown() {
         System.clearProperty("ai.djl.pytorch.native_helper");
         System.clearProperty("LIBSTDCXX_LIBRARY_PATH");
-        System.clearProperty("PYTORCH_PRECXX11");
     }
 
     @Test