[ci] Fixes nightly pytorch jni build (#3196)

deepjavalibrary · May 15, 2024 · c218449 · c218449
1 parent d8fa216
commit c218449
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
@@ -88,6 +88,7 @@ jobs:
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
           CUDA_VERSION=${{ github.event.inputs.cuda }}
+          export CUDA_VERSION=${CUDA_VERSION:-cu121}
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=$CUDA_VERSION -Ppt_version=$PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
       - name: Configure AWS Credentials
@@ -140,6 +141,7 @@ jobs:
           rm -rf ~/.djl.ai
           export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
           CUDA_VERSION=${{ github.event.inputs.cuda }}
+          export CUDA_VERSION=${CUDA_VERSION:-cu121}
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=$CUDA_VERSION -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v2