Repository: deepjavalibrary/djl
Pull request #2866 (2023-11-20 23:12:14)

[CI] Fixes flaky early stopping test
This changes the test from requiring the early stopping at a particular epoch to ensuring that it stops early due to time constraints. This should help reduce the chances of being flaky due to race conditions.

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1365 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`6c27157`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.24%.
> Report is 917 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...va/ai/djl/modality/nlp/generate/TextGenerator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9UZXh0R2VuZXJhdG9yLmphdmE=) | 2.81% | [276 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/nlp/generate/SeqBatcher.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaGVyLmphdmE=) | 0.75% | [132 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...ity/nlp/generate/ContrastiveSeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9Db250cmFzdGl2ZVNlcUJhdGNoU2NoZWR1bGVyLmphdmE=) | 2.97% | [98 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/djl/modality/nlp/generate/SeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaFNjaGVkdWxlci5qYXZh) | 9.83% | [55 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/cv/BufferedImageFactory.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0J1ZmZlcmVkSW1hZ2VGYWN0b3J5LmphdmE=) | 40.96% | [47 Missing and 2 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...a/ai/djl/modality/nlp/generate/StepGeneration.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TdGVwR2VuZXJhdGlvbi5qYXZh) | 2.04% | [48 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [api/src/main/java/ai/djl/ndarray/NDArray.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheS5qYXZh) | 43.42% | [39 Missing and 4 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...n/java/ai/djl/modality/cv/output/CategoryMask.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9DYXRlZ29yeU1hc2suamF2YQ==) | 22.00% | [39 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheUFkYXB0ZXIuamF2YQ==) | 71.21% | [31 Missing and 7 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | 37.50% | [35 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| ... and [76 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2866      +/-   ##
============================================
+ Coverage     72.08%   72.24%   +0.15%     
- Complexity     5126     7161    +2035     
============================================
  Files           473      708     +235     
  Lines         21970    31915    +9945     
  Branches       2351     3314     +963     
============================================
+ Hits          15838    23056    +7218     
- Misses         4925     7285    +2360     
- Partials       1207     1574     +367     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2866?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index d09d144ce..1c46a4204 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -136,7 +136,7 @@ jobs:
           echo $PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
-          ./gradlew :engines:pytorch:pytorch-native:cleanJNI clean
+          ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           rm -rf ~/.djl.ai
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials
diff --git a/.github/workflows/nightly_publish.yml b/.github/workflows/nightly_publish.yml
index 64ac23e85..58a90f45b 100644
--- a/.github/workflows/nightly_publish.yml
+++ b/.github/workflows/nightly_publish.yml
@@ -122,11 +122,11 @@ jobs:
           ./gradlew :integration:test "-Dai.djl.default_engine=OnnxRuntime"
           ./gradlew :integration:clean
 
-  test-cuda-118:
+  test-cuda-121:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: [ self-hosted, gpu ]
     container:
-      image: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu18.04
+      image: nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu20.04
       options: --gpus all --runtime=nvidia
     timeout-minutes: 30
     needs: create-runners
@@ -163,7 +163,7 @@ jobs:
   publish:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: ubuntu-latest
-    needs: [ build, test-pytorch, test-tensorflow, test-aarch64, test-cuda-118 ]
+    needs: [ build, test-pytorch, test-tensorflow, test-aarch64, test-cuda-121 ]
     steps:
       - uses: actions/checkout@v3
       - name: Set up JDK 11
@@ -181,10 +181,9 @@ jobs:
       - name: Publish to snapshot repository
         if: ${{ github.event.inputs.mode == '' || github.event.inputs.mode == 'snapshot' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.11.0 -Psnapshot
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.12.1 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.13.1 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.0.1 -Psnapshot
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.1.1 -Psnapshot
           ./gradlew clean engines:ml:xgboost:publish -Pgpu -Psnapshot
           ./gradlew clean publish -Psnapshot
           cd bom
@@ -197,10 +196,9 @@ jobs:
       - name: Publish to staging repository
         if: ${{ github.event.inputs.mode == 'staging' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.11.0 -P${{ github.event.inputs.mode }}
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.12.1 -P${{ github.event.inputs.mode }}
           ./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.0.1 -P${{ github.event.inputs.mode }}
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=2.1.1 -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
@@ -246,7 +244,7 @@ jobs:
   stop-runners:
     if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
     runs-on: [ self-hosted, scheduler ]
-    needs: [ create-runners, test-aarch64, test-cuda-118 ]
+    needs: [ create-runners, test-aarch64, test-cuda-121 ]
     steps:
       - name: Stop all instances
         run: |
diff --git a/bom/build.gradle b/bom/build.gradle
index 4708978b5..736e8e4a5 100644
--- a/bom/build.gradle
+++ b/bom/build.gradle
@@ -109,21 +109,21 @@ publishing {
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-mkl", "linux-x86_64", "${mxnet_version}")
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-mkl", "win-x86_64", "${mxnet_version}")
                     addDependency(dependencies, "ai.djl.mxnet", "mxnet-native-cu112mkl", "linux-x86_64", "${mxnet_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-aarch64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "linux-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "win-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-aarch64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu116", "linux-x86_64", "1.12.1")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu116", "win-x86_64", "1.12.1")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu116-precxx11", "linux-x86_64", "1.12.1")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "osx-aarch64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu", "win-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cpu-precxx11", "linux-aarch64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121", "win-x86_64", "${pytorch_version}-SNAPSHOT")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu121-precxx11", "linux-x86_64", "${pytorch_version}-SNAPSHOT")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117", "linux-x86_64", "1.13.1")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117", "win-x86_64", "1.13.1")
                     addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu117-precxx11", "linux-x86_64", "1.13.1")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "linux-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "win-x86_64", "${pytorch_version}")
-                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118-precxx11", "linux-x86_64", "${pytorch_version}")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "linux-x86_64", "2.0.1")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118", "win-x86_64", "2.0.1")
+                    addDependency(dependencies, "ai.djl.pytorch", "pytorch-native-cu118-precxx11", "linux-x86_64", "2.0.1")
                     addDependency(dependencies, "ai.djl.tensorflow", "tensorflow-native-cpu", "osx-x86_64", "${tensorflow_version}")
                     addDependency(dependencies, "ai.djl.tensorflow", "tensorflow-native-cpu", "linux-x86_64", "${tensorflow_version}")
                     addDependency(dependencies, "ai.djl.tensorflow", "tensorflow-native-cpu", "win-x86_64", "${tensorflow_version}")
diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
index fbe0df551..d3451ef25 100644
--- a/engines/pytorch/pytorch-engine/README.md
+++ b/engines/pytorch/pytorch-engine/README.md
@@ -46,6 +46,8 @@ The following table illustrates which pytorch version that DJL supports:
 
 | PyTorch engine version | PyTorch native library version            |
 |------------------------|-------------------------------------------|
+| 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         |
 | pytorch-engine:0.24.0  | 1.11.0, 1.12.1, **1.13.1**, 2.0.1         |
 | pytorch-engine:0.23.0  | 1.11.0, 1.12.1, **1.13.1**, 2.0.1         |
 | pytorch-engine:0.22.1  | 1.11.0, 1.12.1, **1.13.1**, 2.0.0         |
diff --git a/engines/pytorch/pytorch-jni/build.gradle b/engines/pytorch/pytorch-jni/build.gradle
index 450c832e8..c2b0ee9dc 100644
--- a/engines/pytorch/pytorch-jni/build.gradle
+++ b/engines/pytorch/pytorch-jni/build.gradle
@@ -24,7 +24,13 @@ processResources {
                 "osx-x86_64/cpu/libdjl_torch.dylib",
                 "win-x86_64/cpu/djl_torch.dll"
         ]
-        if (ptVersion.startsWith("2.0.")) {
+        if (ptVersion.startsWith("2.1.")) {
+            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
+            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")
+        } 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")
diff --git a/engines/pytorch/pytorch-native/CMakeLists.txt b/engines/pytorch/pytorch-native/CMakeLists.txt
index 4453186be..4eddbfc35 100644
--- a/engines/pytorch/pytorch-native/CMakeLists.txt
+++ b/engines/pytorch/pytorch-native/CMakeLists.txt
@@ -64,7 +64,7 @@ add_library(djl_torch SHARED ${SOURCE_FILES})
 if(NOT BUILD_ANDROID)
     target_link_libraries(djl_torch "${TORCH_LIBRARIES}")
     target_include_directories(djl_torch PUBLIC build/include ${JNI_INCLUDE_DIRS} ${UTILS_INCLUDE_DIR})
-    set_property(TARGET djl_torch PROPERTY CXX_STANDARD 14)
+    set_property(TARGET djl_torch PROPERTY CXX_STANDARD 17)
     # We have to kill the default rpath and use current dir
     set(CMAKE_SKIP_RPATH TRUE)
     if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")
diff --git a/engines/pytorch/pytorch-native/build.gradle b/engines/pytorch/pytorch-native/build.gradle
index b4a195e10..99a658bf3 100644
--- a/engines/pytorch/pytorch-native/build.gradle
+++ b/engines/pytorch/pytorch-native/build.gradle
@@ -24,6 +24,8 @@ if (project.hasProperty("cu11")) {
         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}")
     }
@@ -88,15 +90,17 @@ def prepareNativeLib(String binaryRoot, String ver) {
 
     def officialPytorchUrl = "https://download.pytorch.org/libtorch"
     def aarch64PytorchUrl = "https://djl-ai.s3.amazonaws.com/publish/pytorch"
-    String cu11
+    String cuda
     if (ver.startsWith("1.11.")) {
-        cu11 = "cu113"
+        cuda = "cu113"
     } else if (ver.startsWith("1.12.")) {
-        cu11 = "cu116"
+        cuda = "cu116"
     } else if (ver.startsWith("1.13.")) {
-        cu11 = "cu117"
+        cuda = "cu117"
     } else if (ver.startsWith("2.0.")) {
-        cu11 = "cu118"
+        cuda = "cu118"
+    } else if (ver.startsWith("2.1.")) {
+        cuda = "cu121"
     } else {
         throw new GradleException("Unsupported PyTorch version: ${ver}")
     }
@@ -105,10 +109,10 @@ def prepareNativeLib(String binaryRoot, String ver) {
             "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-win-shared-with-deps-${ver}%2Bcpu.zip"              : "cpu/win-x86_64",
-            "${cu11}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cu11}.zip": "${cu11}/linux-x86_64",
-            "${cu11}/libtorch-win-shared-with-deps-${ver}%2B${cu11}.zip"      : "${cu11}/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",
             "cpu/libtorch-shared-with-deps-${ver}%2Bcpu.zip"                  : "cpu-precxx11/linux-x86_64",
-            "${cu11}/libtorch-shared-with-deps-${ver}%2B${cu11}.zip"          : "${cu11}-precxx11/linux-x86_64"
+            "${cuda}/libtorch-shared-with-deps-${ver}%2B${cuda}.zip"          : "${cuda}-precxx11/linux-x86_64"
     ]
 
     def aarch64Files = [
@@ -138,17 +142,12 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
             from zipTree(file)
             into outputDir
         }
-        // CPU dependencies
-        copy {
-            from("${outputDir}/libtorch/lib/") {
-                include "libc10.*", "c10.dll", "libiomp5*.*", "libarm_compute*.*", "libgomp*.*", "libnvfuser_codegen.so", "libtorch.*", "libtorch_cpu.*", "torch.dll", "torch_cpu.dll", "fbgemm.dll", "asmjit.dll", "uv.dll", "nvfuser_codegen.dll"
-            }
-            into("${outputDir}/native/lib")
-        }
-        // GPU dependencies
+        delete "${outputDir}/libtorch/lib/*.lib"
+        delete "${outputDir}/libtorch/lib/*.a"
+
         copy {
             from("${outputDir}/libtorch/lib/") {
-                include "libtorch_cuda*.so", "torch_cuda*.dll", "libc10_cuda.so", "c10_cuda.dll", "libcaffe2_nvrtc.so", "libnvrtc*.so.*", "libcudart*.*", "*nvToolsExt*.*", "cudnn*.dll", "caffe2_nvrtc.dll", "nvrtc64*.dll", "uv.dll", "libcublas*", "zlibwapi.dll"
+                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"
             }
             into("${outputDir}/native/lib")
         }
@@ -287,9 +286,9 @@ tasks.register('uploadS3') {
                 "${BINARY_ROOT}/cpu/win-x86_64/native/lib/",
                 "${BINARY_ROOT}/cpu-precxx11/linux-aarch64/native/lib/",
                 "${BINARY_ROOT}/cpu-precxx11/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu118/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu118/win-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu118-precxx11/linux-x86_64/native/lib/"
+                "${BINARY_ROOT}/cu121/linux-x86_64/native/lib/",
+                "${BINARY_ROOT}/cu121/win-x86_64/native/lib/",
+                "${BINARY_ROOT}/cu121-precxx11/linux-x86_64/native/lib/"
         ]
         uploadDirs.each { item ->
             fileTree(item).files.name.each {
diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
index 78c59d6bf..b3cc2f117 100755
--- a/engines/pytorch/pytorch-native/build.sh
+++ b/engines/pytorch/pytorch-native/build.sh
@@ -23,7 +23,7 @@ ARCH=$4
 
 if [[ ! -d "libtorch" ]]; then
   if [[ $PLATFORM == 'linux' ]]; then
-    if [[ ! "$FLAVOR" =~ ^(cpu|cu102|cu113|cu116|cu117|cu118)$ ]]; then
+    if [[ ! "$FLAVOR" =~ ^(cpu|cu102|cu113|cu116|cu117|cu118|cu121)$ ]]; then
       echo "$FLAVOR is not supported."
       exit 1
     fi
@@ -62,6 +62,12 @@ mkdir classes
 javac -sourcepath ../../pytorch-engine/src/main/java/ ../../pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java -h include -d classes
 cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_VERSION=${PT_VERSION} -DUSE_CUDA=$USE_CUDA ..
 cmake --build . --config Release -- -j "${NUM_PROC}"
+if [[ "$FLAVOR" = cu* ]]; then
+  # avoid link with libcudart.so.11.0
+  sed -i -r "s/\/usr\/local\/cuda(.{5})?\/lib64\/lib(cudart|nvrtc).so//g" CMakeFiles/djl_torch.dir/link.txt
+  rm libdjl_torch.so
+  . CMakeFiles/djl_torch.dir/link.txt
+fi
 
 if [[ $PLATFORM == 'darwin' ]]; then
   install_name_tool -add_rpath @loader_path libdjl_torch.dylib
diff --git a/gradle.properties b/gradle.properties
index 9bd39ebbf..20426a4d0 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -13,7 +13,7 @@ systemProp.org.gradle.internal.publish.checksums.insecure=true
 
 djl_version=0.26.0
 mxnet_version=1.9.1
-pytorch_version=2.0.1
+pytorch_version=2.1.1
 tensorflow_version=2.10.1
 tflite_version=2.6.2
 trt_version=8.4.1


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
index c3dd34d03..426a7fea7 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
@@ -174,11 +174,12 @@ public class EarlyStoppingListenerTest {
                 trainer.setMetrics(metrics);
 
                 try {
-                    // Set epoch to 5 as we expect the early stopping to stop after the second epoch
-                    EasyTrain.fit(trainer, 5, trainMnistDataset, testMnistDataset);
+                    // Set epoch to 10 as we expect the early stopping to stop after the second
+                    // epoch
+                    EasyTrain.fit(trainer, 10, trainMnistDataset, testMnistDataset);
                 } catch (EarlyStoppingListener.EarlyStoppedException e) {
                     Assert.assertTrue(e.getMessage().contains("ms elapsed >= 1 maxMillis"));
-                    Assert.assertEquals(e.getStopEpoch(), 1);
+                    Assert.assertTrue(e.getStopEpoch() < 10); // Stop epoch is before 10
                 }
 
                 TrainingResult trainingResult = trainer.getTrainingResult();