mirror of
https://github.com/hanzo-ml/katib.git
synced 2026-09-24 00:31:27 +00:00
20 lines
683 B
Docker
20 lines
683 B
Docker
# We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
|
|
# PyTorch=2.2.0, cuda=12.3.2
|
|
# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html#rel-24-01
|
|
FROM nvcr.io/nvidia/pytorch:24.01-py3
|
|
|
|
ADD examples/v1beta1/trial-images/pytorch-mnist /opt/pytorch-mnist
|
|
|
|
WORKDIR /opt/pytorch-mnist
|
|
|
|
# Add folder for the logs.
|
|
RUN mkdir /katib
|
|
RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
|
|
|
|
RUN chgrp -R 0 /opt/pytorch-mnist \
|
|
&& chmod -R g+rwX /opt/pytorch-mnist \
|
|
&& chgrp -R 0 /katib \
|
|
&& chmod -R g+rwX /katib
|
|
|
|
ENTRYPOINT ["python3", "/opt/pytorch-mnist/mnist.py"]
|