Update Dockerfile

Dockerfile

@@ -15,8 +15,8 @@ RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86
     dpkg -i cuda-keyring_1.0-1_all.deb && \
     apt-get update
 
-# Install CUDA Toolkit (includes libcublas)
-RUN apt-get install -y cuda-toolkit-11-7 libcudnn8
+# Install CUDA Toolkit and cuDNN
+RUN apt-get install -y cuda-toolkit-11-7 libcudnn8 libcudnn8-dev
 
 # Create a user
 RUN useradd -m -u 1000 user
@@ -36,9 +36,11 @@ ENV HOME=/home/user \
     CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda \
     CUDA_HOME=/usr/local/cuda \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
-    CMAKE_PREFIX_PATH=/usr/local/cuda \
     TORCH_CUDA_ARCH_LIST="8.6"
 
+# Set CUDA & cuDNN paths
+ENV CMAKE_PREFIX_PATH="/usr/local/cuda:/usr/lib/x86_64-linux-gnu/"
+
 # Set working directory
 WORKDIR $HOME/app