Merge pull request #669 from OptimalScale/rpan-support-a6000

Add flash attention install for A6000
OptimalScale · Oct 29, 2023 · f846912 · f846912
2 parents 2a311b3 + 7932633
commit f846912
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