# Expecting base image to be the Deep Learning Container image built by ../cu117/Dockerfile.ec2.gpu
ARG BASE_IMAGE="../cu118/Dockerfile.gpu"

FROM $BASE_IMAGE

RUN pip install tensorboardX==2.4

# Add any script or repo as required
COPY mnist.py /var/mnist.py

ENTRYPOINT ["python", "/var/mnist.py"]