update release_build.yml

pytorch · Oct 4, 2023 · 2bde575 · 2bde575
1 parent e8c7b33
commit 2bde575
Showing 1 changed file with 23 additions and 7 deletions.
diff --git a/.github/workflows/release_build.yml b/.github/workflows/release_build.yml
@@ -25,15 +25,19 @@ jobs:
          - os: linux.2xlarge
            python-version: 3.8
            python-tag: "py38"
-           cuda-tag: "cu11"
+           cuda-tag: "cu118"
          - os: linux.2xlarge
            python-version: 3.9
            python-tag: "py39"
-           cuda-tag: "cu11"
+           cuda-tag: "cu118"
          - os: linux.2xlarge
            python-version: '3.10'
            python-tag: "py310"
-           cuda-tag: "cu11"
+           cuda-tag: "cu118"
+         - os: linux.2xlarge
+           python-version: '3.11'
+           python-tag: "py311"
+           cuda-tag: "cu118"
     steps:
     # Checkout the repository to the GitHub Actions runner
     - name: Check ldd --version
@@ -69,7 +73,11 @@ jobs:
     - name: Install PyTorch and CUDA
       shell: bash
       run: |
-        conda install -n build_binary -y pytorch pytorch-cuda=11.7 -c pytorch-test -c nvidia
+        conda install -n build_binary pip install torch --index-url https://download.pytorch.org/whl/test/cu118
+    - name: Install fbgemm
+      shell: bash
+      run: |
+        conda install -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/test/cu118
     - name: Install Dependencies
       shell: bash
       run: |
@@ -102,8 +110,8 @@ jobs:
     strategy:
       matrix:
         os: [linux.4xlarge.nvidia.gpu]
-        python-version: [3.8, 3.9, "3.10"]
-        cuda-tag: ["cu11"]
+        python-version: [3.8, 3.9, "3.10", "3.11"]
+        cuda-tag: ["cu118"]
     needs: build_on_cpu
     # the glibc version should match the version of the one we used to build the binary
     # for this case, it's 2.26
@@ -168,8 +176,16 @@ jobs:
     - name: Install PyTorch and CUDA
       shell: bash
       run: |
-        conda install -n build_binary -y pytorch pytorch-cuda=11.7 -c pytorch-test -c nvidia
+        conda install -n build_binary pip install torch --index-url https://download.pytorch.org/whl/test/cu118
     # download wheel from GHA
+    - name: Install fbgemm
+      shell: bash
+      run: |
+        conda install -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/test/cu118
+    - name: Install torchmetrics
+      shell: bash
+      run: |
+        conda install -n build_binary pip install torchmetrics==1.0.3
     - name: Download wheel
       uses: actions/download-artifact@v2
       with: