robocar-training/Dockerfile.gpu

16 lines
374 B
Docker

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