From c99e625045439803f710e3c2cd502e3a0ddf0a8d Mon Sep 17 00:00:00 2001 From: rohithkrn Date: Mon, 6 Nov 2023 09:16:29 -0800 Subject: [PATCH] install datasets in trtllm container (#1270) --- serving/docker/tensorrt-llm.Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/serving/docker/tensorrt-llm.Dockerfile b/serving/docker/tensorrt-llm.Dockerfile index 1d8d34eb877..6495a999ed5 100644 --- a/serving/docker/tensorrt-llm.Dockerfile +++ b/serving/docker/tensorrt-llm.Dockerfile @@ -66,7 +66,7 @@ RUN apt-get update && apt-get install -y wget unzip openmpi-bin libopenmpi-dev l # Install PyTorch and TensorRT RUN pip install torch==${TORCH_VERSION} transformers==${transformers_version} accelerate==${accelerate_version} ${peft_wheel} sentencepiece \ - mpi4py cuda-python==${cuda_python_version} onnx polygraphy && \ + mpi4py cuda-python==${cuda_python_version} onnx polygraphy datasets && \ pip install --no-cache-dir --extra-index-url https://pypi.nvidia.com tensorrt==${tensorrtlibs_version} && \ pip install --no-deps ${trtllm_wheel} && \ pip3 cache purge