From 593a5634142e52b537f9bfe1bbf131c7ade1185f Mon Sep 17 00:00:00 2001
From: OlivierDehaene <olivier@huggingface.co>
Date: Thu, 27 Apr 2023 19:18:33 +0200
Subject: [PATCH] feat(docker): add nvidia env vars (#255)

---
 Dockerfile | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/Dockerfile b/Dockerfile
index c17a9414..c65d3dee 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -164,6 +164,11 @@ FROM base as sagemaker
 COPY sagemaker-entrypoint.sh entrypoint.sh
 RUN chmod +x entrypoint.sh
 
+# NVIDIA env vars
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
+
 ENTRYPOINT ["./entrypoint.sh"]
 
 # Final image