Update Dockerfile

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.3.2-cudnn9-runtime-ubuntu22.04
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04 AS build
 
 ARG DEBIAN_FRONTEND=noninteractive
 
@@ -23,6 +23,8 @@ COPY ./requirements.txt /app/requirements.txt
 RUN pip3 install --no-cache-dir --upgrade -r /app/requirements.txt
 RUN CMAKE_ARGS="-DLLAMA_CUDA=on" pip3 install --no-cache-dir llama-cpp-python
 
+FROM nvcr.io/nvidia/pytorch:24.02-py3
+COPY --from=build /app /app
 
 RUN useradd -m -u 1000 user
 USER user