updated Dockerfile.gpu to install cuda-toolkit

tensorflow · Feb 6, 2024 · 96a35a5 · 96a35a5
1 parent b23acfc
commit 96a35a5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorflow_serving/tools/docker/Dockerfile.gpu b/tensorflow_serving/tools/docker/Dockerfile.gpu
@@ -60,6 +60,8 @@ RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machi
     && apt-get clean && \
     rm -rf /var/lib/apt/lists/*;
 
+RUN apt-get update && apt-get install -y cuda-toolkit-11-8
+
 # Install TF Serving GPU pkg
 COPY --from=build_image /usr/local/bin/tensorflow_model_server /usr/bin/tensorflow_model_server