Include cuda 12.8

This commit is contained in:
Jake Poznanski 2025-06-02 22:52:28 +00:00
parent 5c524b53ac
commit 5c52e016e6
2 changed files with 2 additions and 2 deletions

View File

@ -61,7 +61,7 @@ COPY olmocr/version.py olmocr/version.py
# Needed to resolve setuptools dependencies # Needed to resolve setuptools dependencies
ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_INDEX_STRATEGY="unsafe-best-match"
RUN uv pip install --system --no-cache -e . RUN uv pip install --system --no-cache -e . --extra-index-url https://download.pytorch.org/whl/cu128
RUN uv pip install --system --no-cache ".[gpu]" --extra-index-url https://download.pytorch.org/whl/cu128 RUN uv pip install --system --no-cache ".[gpu]" --extra-index-url https://download.pytorch.org/whl/cu128
RUN uv pip install --system https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl RUN uv pip install --system https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl
RUN uv pip install --system --no-cache ".[bench]" RUN uv pip install --system --no-cache ".[bench]"

View File

@ -36,7 +36,7 @@ dependencies = [
"zstandard", "zstandard",
"boto3", "boto3",
"httpx", "httpx",
"torch>=2.5.1", "torch>=2.7.0",
"transformers>=4.51.1", "transformers>=4.51.1",
"img2pdf", "img2pdf",
"beaker-py", "beaker-py",