diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt
index 9bcf7b0..c91dc8e 100644
--- a/docker/Dockerfile.tensorrt
+++ b/docker/Dockerfile.tensorrt
@@ -1,8 +1,14 @@
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
 ARG BASE_TAG=23.10-py3
+ARG CUDA_ARCHITECTURES="89-real"
 
 FROM ${BASE_IMAGE}:${BASE_TAG} as base
 
 WORKDIR /home
 COPY script/install_trt_llm.sh install_trt_llm.sh
-RUN bash install_trt_llm.sh && rm install_trt_llm.sh
\ No newline at end of file
+RUN bash install_trt_llm.sh ${CUDA_ARCHITECTURES} && rm install_trt_llm.sh
+
+COPY scripts/setup.sh /home
+RUN apt update && bash setup.sh && rm setup.sh
+COPY requirements/ /home
+RUN pip install -r server.txt && rm server.txt