feat(docker): add nvidia env vars (#255)
This commit is contained in:
parent
f092ba9b22
commit
593a563414
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue