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