Add flash attention install for A6000

OptimalScale · Oct 29, 2023 · 7932633 · 7932633
1 parent 647afd8
commit 7932633
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/install.sh b/install.sh
@@ -3,6 +3,6 @@
 pip install -e .
 
 gpu_state="$(nvidia-smi --query-gpu=name --format=csv,noheader)"
-if [[ "${gpu_state}" == *"A100"* || "${gpu_state}" == *"A40"* ]]; then
+if [[ "${gpu_state}" == *"A100"* || "${gpu_state}" == *"A40"* || "${gpu_state}" == *"A6000"* ]]; then
   pip install flash-attn==2.0.2
-fi
+fi