FROM gooseai/torch-base:6cfdc11
RUN apt-get install -y cuda-nvcc-11-3 cuda-nvml-dev-11-3 libcurand-dev-11-3 \
                       libcublas-dev-11-3 libcusparse-dev-11-3 \
                       libcusolver-dev-11-3 cuda-nvprof-11-3 \
                       ninja-build && \
    apt-get clean

RUN mkdir -p /app
WORKDIR /app
# Requirements
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt
# Then our source
COPY . .
CMD [ "/usr/bin/python3", "hypertrain.py" ]
