Update cuda images (#246)

huggingface · Aug 19, 2024 · bca546a · bca546a
1 parent 216172c
commit bca546a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/.github/workflows/test_cli_cuda_onnxruntime.yaml b/.github/workflows/test_cli_cuda_onnxruntime.yaml
@@ -26,7 +26,7 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
 
 env:
-  IMAGE: ghcr.io/huggingface/optimum-benchmark:latest-cuda-ort
+  IMAGE: ghcr.io/huggingface/optimum-benchmark:latest-cuda
 
 jobs:
   run_cli_cuda_onnxruntime_tests:

diff --git a/docker/cuda/Dockerfile b/docker/cuda/Dockerfile
@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG CUDNN_VERSION=8
-ARG CUDA_VERSION=12.1.1
+ARG CUDA_VERSION=12.4.1
 ARG UBUNTU_VERSION=22.04
 
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel-ubuntu${UBUNTU_VERSION}
+FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION}
 
 # Install necessary packages
 ENV DEBIAN_FRONTEND noninteractive
@@ -28,7 +27,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     pip install --no-cache-dir --upgrade pip setuptools wheel 
 
 # Install PyTorch
-ARG TORCH_CUDA=cu121
+ARG TORCH_CUDA=cu124
 ARG TORCH_VERSION=stable
 
 RUN if [ "${TORCH_VERSION}" = "stable" ]; then \