fix: test ci

.github/workflows/ci.yml

@@ -154,12 +154,19 @@ jobs:
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Install CUDA
+        run: |
+          bash .github/workflows/cuda/cu113-Linux.sh
+          source .github/workflows/cuda/cu113-Linux-env.sh
+          nvcc --version
+        shell:
+          bash
       - name: Prepare enviroment
         run: |
           python -m pip install --upgrade pip
           python -m pip install wheel pytest pytest-cov nvidia-pyindex
           {
-            python -m pip install torch==1.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
+            pip install torch==1.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
             python -m pip install git+https://github.com/HazyResearch/flash-attention.git
           } || {
             echo "flash attention was not installed."

.github/workflows/cuda/cu113-Linux-env.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+CUDA_HOME=/usr/local/cuda-11.3
+LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
+PATH=${CUDA_HOME}/bin:${PATH}
+
+export FORCE_CUDA=1
+export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
+export CUDA_HOME=/usr/local/cuda-11.3

.github/workflows/cuda/cu113-Linux.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+OS=ubuntu1804
+
+wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin
+sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600
+wget -nv https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda-repo-${OS}-11-3-local_11.3.0-465.19.01-1_amd64.deb
+sudo dpkg -i cuda-repo-${OS}-11-3-local_11.3.0-465.19.01-1_amd64.deb
+sudo apt-key add /var/cuda-repo-${OS}-11-3-local/7fa2af80.pub
+
+sudo apt-get -qq update
+sudo apt install cuda cuda-nvcc-11-3 cuda-libraries-dev-11-3
+sudo apt clean
+
+rm -f https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda-repo-${OS}-11-3-local_11.3.0-465.19.01-1_amd64.deb