From a6a6759e36ce0663a07efb9ce15088edf9da0d36 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Mon, 19 Aug 2024 10:15:48 +0200
Subject: [PATCH] update cuda images

---
 .github/workflows/test_cli_cuda_onnxruntime.yaml | 2 +-
 docker/cuda/Dockerfile                           | 7 +++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/test_cli_cuda_onnxruntime.yaml b/.github/workflows/test_cli_cuda_onnxruntime.yaml
index 355296dc..8bdabd87 100644
--- 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
index 6ec05d28..61425635 100644
--- 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 \