From 116cb0ccdb10cfb3e3bf2246363342cb7534c019 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Tue, 14 May 2024 23:30:27 -0700 Subject: [PATCH] [ci] Fixes nightly pytorch jni build --- .github/workflows/native_jni_s3_pytorch.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml index 18aab758461..95fd4a8cbac 100644 --- 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