FROM tensorflow/tensorflow:2.3.1-gpu

# Install pip requirements
COPY requirements.txt /tmp/
RUN pip3 install -r /tmp/requirements.txt
COPY . /tmp/

# Set the working directory.
WORKDIR /home/sparsegradientflow

# Copy the source code
COPY * ./