Pg hard code cudnn tensorrt version (#10)

* chore: switch to ubuntu nvidia packages

* fix: ARG fix CuDNN & TensorRT

* fix: dynamical set cuda version TensorRT

Author: PGerardi

Dockerfile

@@ -1,12 +1,11 @@
 # syntax=docker/dockerfile:1
 ARG PYTHON_VERSION=3.8
+FROM python:$PYTHON_VERSION-slim AS base
+
 ARG CUDA_VERSION=11.8
 ARG CUDNN_VERSION=8.9.4.25
 ARG TENSORRT_VERSION=8.6.1.6
-FROM python:$PYTHON_VERSION-slim AS base
-ENV CUDA_VERSION=11.8
-ENV CUDNN_VERSION=8.9.4.25
-ENV TENSORRT_VERSION=8.6.1.6
+
 # Install wget.
 RUN apt-get update && \
     apt-get install --yes wget && \
@@ -22,14 +21,20 @@ ENV PATH=/opt/conda/bin:$PATH
 # Install CUDA.
 RUN conda install --channel nvidia --yes cuda-runtime="$CUDA_VERSION"
 
-# Install cuDNN and TensorRT.
-RUN wget --quiet https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.1-1_all.deb && \
+# Cuda version compatible with tensorRT version
+RUN if [ "$CUDA_VERSION" = "12.2" ]; then \
+        export CUDA_TENSORRT_VERSION=12.0; \
+    else \
+        export CUDA_TENSORRT_VERSION=$CUDA_VERSION; \
+    fi \
+    && \
+    # Now use $CUDA_TENSORRT_VERSION in the same RUN command 
+    wget --quiet https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb && \
     dpkg -i cuda-keyring_1.1-1_all.deb && \
     rm cuda-keyring_1.1-1_all.deb && \
     apt-get update && \
     apt-get install --yes \
         libcudnn8=$CUDNN_VERSION-1+cuda$CUDA_VERSION \
-        libnvinfer-lean8=$TENSORRT_VERSION-1+cuda$CUDA_VERSION && \
+        libnvinfer-lean8=$TENSORRT_VERSION-1+cuda$CUDA_TENSORRT_VERSION && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
-