feat: add cuDNN and TensorRT (#7)

Author: lsorber

Dockerfile

@@ -1,6 +1,8 @@
 # syntax=docker/dockerfile:1
 ARG PYTHON_VERSION=3.8
 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
 
 # Install wget.
@@ -17,4 +19,15 @@ ENV PATH=/opt/conda/bin:$PATH
 
 # Install CUDA.
 RUN conda install --channel nvidia --yes cuda-runtime="$CUDA_VERSION"
-RUN conda install -c conda-forge cudnn
+
+# 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 && \
+    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 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+