Add torchtrition dependency for wheels

ROCm · Dec 1, 2022 · 475100b · 475100b
1 parent dd3c250
commit 475100b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/manywheel/build_cuda.sh b/manywheel/build_cuda.sh
@@ -255,6 +255,12 @@ else
     exit 1
 fi
 
+# TODO: Remove me when Triton has a proper release channel
+if [[ $(uname) == "Linux" && -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
+    TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.github/ci_commit_pins/triton.txt)
+    export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="torchtriton==2.0.0+${TRITON_SHORTHASH}"
+fi
+
 # builder/test.sh requires DESIRED_CUDA to know what tests to exclude
 export DESIRED_CUDA="$cuda_version_nodot"