From 3d0a3003e22e34ce1023ab0206a952c139492f2f Mon Sep 17 00:00:00 2001
From: atalman <atalman@fb.com>
Date: Tue, 13 Jun 2023 11:58:33 -0700
Subject: [PATCH] Add pytorch-triton to small wheel

---
 manywheel/build_cuda.sh | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/manywheel/build_cuda.sh b/manywheel/build_cuda.sh
index 3bc3a49f9..8af346e85 100644
--- a/manywheel/build_cuda.sh
+++ b/manywheel/build_cuda.sh
@@ -272,9 +272,14 @@ else
 fi
 
 # TODO: Remove me when Triton has a proper release channel
-if [[ $(uname) == "Linux" && -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
+if [[ $(uname) == "Linux" ]]; then
     TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.github/ci_commit_pins/triton.txt)
-    export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton==2.1.0+${TRITON_SHORTHASH}"
+
+    if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
+        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton==2.1.0+${TRITON_SHORTHASH}"
+    else
+        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton==2.1.0+${TRITON_SHORTHASH}"
+    fi
 fi
 
 # builder/test.sh requires DESIRED_CUDA to know what tests to exclude