Dockerfile

FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

RUN apt-get update && apt-get install -y python3 python3-pip git
RUN ln -s /usr/bin/python3 /usr/bin/python

WORKDIR /app
COPY . .

RUN pip install --upgrade pip
RUN pip install -r requirements.txt

RUN mkdir -p /.cache && chmod a+rwxt /.cache
RUN mkdir -p /.triton && chmod a+rwxt /.triton
RUN mkdir -p /.results && chmod a+rwxt /.results

CMD ["python", "train.py"]