FROM pytorch/pytorch:0.4_cuda9_cudnn7

WORKDIR /workspace
RUN chmod -R a+w /workspace
ADD ./mnist_DDP.py /opt/pytorch_dist_mnist/

ENTRYPOINT ["python", "-u", "/opt/pytorch_dist_mnist/mnist_DDP.py", "--gpu", "--modelpath", "/mnt/kubeflow-gcfs/pytorch/model"]