From ee7ad322ec37d1c81d66cb4a84cd232ba8e89074 Mon Sep 17 00:00:00 2001
From: Emmanuel Benazera <emmanuel.benazera@jolibrain.com>
Date: Sun, 3 Nov 2024 09:29:31 +0000
Subject: [PATCH] chore: cuda 12.1.1 for GPU docker images

---
 docker/gpu.Dockerfile | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/docker/gpu.Dockerfile b/docker/gpu.Dockerfile
index 563e95ebe..0473fdff2 100644
--- a/docker/gpu.Dockerfile
+++ b/docker/gpu.Dockerfile
@@ -1,5 +1,14 @@
 # syntax = docker/dockerfile:1.0-experimental
-FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS build
+
+ARG DD_UBUNTU_VERSION=22.04
+ARG DD_CUDA_VERSION=12.1.1
+ARG DD_CUDNN_VERSION=8
+FROM nvidia/cuda:${DD_CUDA_VERSION}-cudnn${DD_CUDNN_VERSION}-devel-ubuntu${DD_UBUNTU_VERSION} AS build
+
+
+RUN echo UBUNTU_VERSION=${DD_UBUNTU_VERSION} >> /image-info
+RUN echo CUDA_VERSION=${DD_CUDA_VERSION} >> /image-info
+RUN echo CUDNN_VERSION=${DD_CUDNN_VERSION} >> /image-info
 
 ARG DEEPDETECT_RELEASE=OFF
 ARG DEEPDETECT_ARCH=gpu
@@ -106,7 +115,7 @@ RUN --mount=type=cache,target=/ccache/ mkdir build && cd build && ../build.sh
 RUN ./docker/get_libs.sh
 
 # Build final Docker image
-FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 AS runtime
+FROM nvidia/cuda:${DD_CUDA_VERSION}-cudnn${DD_CUDNN_VERSION}-runtime-ubuntu${DD_UBUNTU_VERSION} AS runtime
 
 ARG DEEPDETECT_ARCH=gpu