Fix ROCm build (#342)

casper-hansen · Feb 15, 2024 · 9477ea5 · 9477ea5
1 parent bcaa8a3
commit 9477ea5
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -209,11 +209,9 @@ jobs:
           python -m pip install --upgrade build setuptools wheel
 
           if [[ "${{ matrix.rocm }}" == "5.7.1" ]]; then
-            echo "Using PyTorch nightly"
-            python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm5.7
+            python -m pip install torch==2.2.0 --index-url https://download.pytorch.org/whl/rocm5.7
           elif [[ "${{ matrix.rocm }}" == "5.6.1" ]]; then
-            echo "Using PyTorch stable"
-            python -m pip install torch --index-url https://download.pytorch.org/whl/rocm5.6
+            python -m pip install torch==2.2.0 --index-url https://download.pytorch.org/whl/rocm5.6
           else
             echo Unknown rocm version for python install
             exit 1