CUDA 12.4 update as default for pypi and gpu ci

pytorch · Oct 21, 2024 · 065647c · 065647c
1 parent 5319d87
commit 065647c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/.github/workflows/release_build.yml b/.github/workflows/release_build.yml
@@ -18,19 +18,19 @@ jobs:
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: 3.9
             python-tag: "py39"
-            cuda-tag: "cu121"
+            cuda-tag: "cu124"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.10'
             python-tag: "py310"
-            cuda-tag: "cu121"
+            cuda-tag: "cu124"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.11'
             python-tag: "py311"
-            cuda-tag: "cu121"
+            cuda-tag: "cu124"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.12'
             python-tag: "py312"
-            cuda-tag: "cu121"
+            cuda-tag: "cu124"
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     with:
       runner: ${{ matrix.os }}
@@ -41,9 +41,9 @@ jobs:
         conda info
         python --version
         conda run -n build_binary python --version
-        conda run -n build_binary pip install torch --index-url https://download.pytorch.org/whl/test/cu121
+        conda run -n build_binary pip install torch
         conda run -n build_binary pip install numpy
-        conda run -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/test/cu121
+        conda run -n build_binary pip install fbgemm-gpu
         conda run -n build_binary python -m pip install -r requirements.txt
         conda run -n build_binary python -c "import torch.distributed"
         echo "torch.distributed succeeded"

diff --git a/.github/workflows/unittest_ci.yml b/.github/workflows/unittest_ci.yml
@@ -23,6 +23,10 @@ jobs:
             python-version: 3.9
             python-tag: "py39"
             cuda-tag: "cu121"
+          - os: linux.g5.12xlarge.nvidia.gpu
+            python-version: 3.9
+            python-tag: "py39"
+            cuda-tag: "cu124"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.10'
             python-tag: "py310"
@@ -31,6 +35,10 @@ jobs:
             python-version: '3.10'
             python-tag: "py310"
             cuda-tag: "cu121"
+          - os: linux.g5.12xlarge.nvidia.gpu
+            python-version: '3.10'
+            python-tag: "py310"
+            cuda-tag: "cu124"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.11'
             python-tag: "py311"
@@ -39,6 +47,10 @@ jobs:
             python-version: '3.11'
             python-tag: "py311"
             cuda-tag: "cu121"
+          - os: linux.g5.12xlarge.nvidia.gpu
+            python-version: '3.11'
+            python-tag: "py311"
+            cuda-tag: "cu124"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.12'
             python-tag: "py312"
@@ -47,6 +59,10 @@ jobs:
             python-version: '3.12'
             python-tag: "py312"
             cuda-tag: "cu121"
+          - os: linux.g5.12xlarge.nvidia.gpu
+            python-version: '3.12'
+            python-tag: "py312"
+            cuda-tag: "cu124"
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     with:
       runner: ${{ matrix.os }}