From 6e22744847bbc32464ab7f313985b6e78a5051b9 Mon Sep 17 00:00:00 2001
From: krishung5 <krish@nvidia.com>
Date: Wed, 11 Oct 2023 17:50:19 -0700
Subject: [PATCH] Remove packages that are only needed for testing

---
 build.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/build.py b/build.py
index 61bdddfd44..7d14a26cc4 100755
--- a/build.py
+++ b/build.py
@@ -1349,20 +1349,17 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
     rm -fr  ${TRT_ROOT}/doc ${TRT_ROOT}/onnx_graphsurgeon ${TRT_ROOT}/python && \
     rm -fr ${TRT_ROOT}/samples  ${TRT_ROOT}/targets/${ARCH}-linux-gnu/samples
 
+# Install required packages for TRT-LLM models
+RUN python3 -m pip install --upgrade pip && \
+        pip3 install transformers && \
+        pip3 install torch
+
 # Uninstall unused nvidia packages
 RUN if pip freeze | grep -q "nvidia.*"; then \
         pip freeze | grep "nvidia.*" | xargs pip uninstall -y; \
     fi
 RUN pip cache purge
 
-# Install required packages for example models
-RUN python3 -m pip install --upgrade pip && \
-        pip3 install transformers && \
-        pip3 install torch && \
-        pip3 install tritonclient[all] && \
-        pip3 install pandas && \
-        pip3 install tabulate
-
 ENV LD_LIBRARY_PATH=/usr/local/tensorrt/lib/:/opt/tritonserver/backends/tensorrtllm:$LD_LIBRARY_PATH
 """