squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CI] Install torch 2.3 using index

Browse Source
This commit is contained in:
Tri Dao 2024-01-30 14:32:29 -08:00
parent daf37a9d8a
commit dc72d960a7
2 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
.github/workflows/publish.yml vendored
View File

@ -133,9 +133,13 @@ jobs:
print(max(min(int(env['MATRIX_CUDA_VERSION']), maxv), minv))" \
)
if [[ ${{ matrix.torch-version }} == *"dev"* ]]; then
# --no-deps because we can't install old versions of pytorch-triton
pip install typing-extensions jinja2
pip install --no-cache-dir --no-deps --pre https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}/torch-${{ matrix.torch-version }}%2Bcu${TORCH_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
if [[ ${MATRIX_TORCH_VERSION} == "2.2" ]]; then
# --no-deps because we can't install old versions of pytorch-triton
pip install typing-extensions jinja2
pip install --no-cache-dir --no-deps --pre https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}/torch-${{ matrix.torch-version }}%2Bcu${TORCH_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
else
pip install --no-cache-dir --pre torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}
fi
else
pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
fi

2
flash_attn/__init__.py
View File

@ -1,4 +1,4 @@
__version__ = "2.5.1"
__version__ = "2.5.1.post1"
from flash_attn.flash_attn_interface import (
flash_attn_func,