From 349c35e4511d9c92cdc50158dbfed28d6d028ff2 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Thu, 27 Apr 2023 18:53:07 +0200
Subject: [PATCH] feat(docker): add nvidia env vars

---
 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