FROM docker://nvidia/cuda:12.9.1-cudnn-devel-ubuntu24.04

RUN apt update && apt install -y python3 python3-pip python3-venv git cmake && apt clean

ENV PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu129

RUN python3 -m pip install --break-system-packages tox torch

ENV CUDA_HOME="/usr/local/cuda"
ENV Torch_DIR=/usr/local/lib/python3.12/dist-packages/torch/share/cmake/Torch/
