FROM docker.io/tensorflow/tensorflow:2.6.0-gpu

COPY requirements.txt .
RUN pip3 install --upgrade pip==20.0.2 && pip3 list && pip3 install -r requirements.txt \
    && pip3 list

WORKDIR /root

# copy the training script inside the container
COPY src/tf_container/train.py /opt/ml/code/train.py

# define train.py as the script entry point
ENV SAGEMAKER_PROGRAM train.py