ci(cuda): install CUDA toolkit to build extension

huggingface · Apr 10, 2024 · c4788b1 · c4788b1
1 parent 8ecc1d3
commit c4788b1
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/.github/workflows/linux-cuda-tests.yml b/.github/workflows/linux-cuda-tests.yml
@@ -27,15 +27,14 @@ jobs:
     runs-on: [self-hosted, single-gpu , nvidia-gpu, a10, ci]
     strategy:
       fail-fast: false
-      matrix:
-        python-version: ["3.11"]
+    container:
+      image: pytorch/pytorch:2.2.1-cuda12.1-cudnn8-devel
 
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
-        with:
-          python-version: ${{ matrix.python-version }}
+      - name: Check CUDA installation
+        run: |
+          nvcc -V
 
       # Install dependencies
       - name: Configure and install dependencies