From 78e0c7eec92f20daa82482ac813b7f066d7041eb Mon Sep 17 00:00:00 2001 From: Kevin Lin <42618777+kevin314@users.noreply.github.com> Date: Wed, 18 Jun 2025 20:02:50 -0500 Subject: [PATCH] Specify cu128 Pytorch installation (#530) Co-authored-by: Edenzzzz Co-authored-by: Wenxuan Tan --- docker/Dockerfile.python3.12 | 4 ++-- pyproject.toml | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index 002e8ed3b..2c42853c5 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -58,14 +58,14 @@ RUN source $HOME/.local/bin/env && \ # Install STA (Sliding Tile Attention) RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/ && \ + cd csrc/attn && \ git submodule update --init --recursive && \ python setup_sta.py install # Install VSA RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/ && \ + cd csrc/attn && \ git submodule update --init --recursive && \ python setup_vsa.py install diff --git a/pyproject.toml b/pyproject.toml index 65ce83ac0..75506e14f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,9 @@ dependencies = [ "av", ] +[tool.uv] +extra-index-url = ["https://download.pytorch.org/whl/cu128"] + [project.optional-dependencies] # flash-attn: pip install flash-attn==2.7.4.post1 --no-cache-dir --no-build-isolation