FROM pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime@sha256:417bd75df6365104c283ea4c1651fb3530d9eb5a4c2fafa51943cff2a94e6385
RUN apt-get update && apt-get install -y --no-install-recommends git tmux asciinema && rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY environment/requirements-lock.txt /app/requirements-lock.txt
RUN python -m pip install --no-cache-dir --no-deps -r /app/requirements-lock.txt
COPY environment/train.py /app/train.py
COPY environment/gputools/ /app/gputools/
COPY environment/data/ /app/data/
RUN mkdir -p /app/output
ENV PYTHONPATH=/app CUBLAS_WORKSPACE_CONFIG=:4096:8
CMD ["python3", "/app/train.py", "--data", "/app/data/public_data/public_data.npz", "--output", "/app/output"]
