[CI] Fix MATRIX_CUDA_VERSION check

Dao-AILab · Aug 14, 2023 · f8dccfc · f8dccfc
1 parent 9c531bd
commit f8dccfc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
@@ -142,7 +142,7 @@ jobs:
           export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
           export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
           # Currently for this setting the runner goes OOM if we pass --threads 4 to nvcc
-          if [[ ${MATRIX_CUDA_VERSION} == "12.1" && ${MATRIX_TORCH_VERSION} == "2.1" ]]; then
+          if [[ ${MATRIX_CUDA_VERSION} == "121" && ${MATRIX_TORCH_VERSION} == "2.1" ]]; then
             export FLASH_ATTENTION_FORCE_SINGLE_THREAD="TRUE"
           fi
           # Limit MAX_JOBS otherwise the github runner goes OOM

diff --git a/flash_attn/__init__.py b/flash_attn/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.0.6.post1"
+__version__ = "2.0.6.post2"
 
 from flash_attn.flash_attn_interface import flash_attn_func
 from flash_attn.flash_attn_interface import flash_attn_kvpacked_func