From eaed605c14e54d34761b56244339a5bdcce37c5d Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Thu, 20 Mar 2025 21:32:27 -0700
Subject: [PATCH 1/9] Upload PyTorch MacOS wheel to S3

---
 .ci/scripts/utils.sh            | 39 +++++++++++++++++++++++----------
 .github/workflows/_unittest.yml |  2 ++
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index ed9f6f21a5b..d67faaa4b21 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -55,17 +55,34 @@ install_pytorch_and_domains() {
   TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
   popd || return
 
-  git clone https://github.com/pytorch/pytorch.git
-
-  # Fetch the target commit
-  pushd pytorch || return
-  git checkout "${TORCH_VERSION}"
-  git submodule update --init --recursive
-
-  export USE_DISTRIBUTED=1
-  # Then build and install PyTorch
-  python setup.py bdist_wheel
-  pip install "$(echo dist/*.whl)"
+  OS_VERSION=$(uname)
+  PYTHON_VERSION=$(python -c 'import platform; print(platform.python_version())')
+  PYTORCH_WHEEL_S3_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${OS_VERSION}/${PYTHON_VERSION}"
+
+  if [[ "${FOUND_PYTORCH_WHEEL:-0}" == "0" ]]; then
+    # Found no such wheel, we will build it from source then
+    git clone https://github.com/pytorch/pytorch.git
+
+    # Fetch the target commit
+    pushd pytorch || return
+    git checkout "${TORCH_VERSION}"
+    git submodule update --init --recursive
+
+    export USE_DISTRIBUTED=1
+    # Then build and install PyTorch
+    python setup.py bdist_wheel
+    pip install "$(echo dist/*.whl)"
+
+    # Only AWS runners have access to S3
+    if command -v aws && [[ -z "${GITHUB_RUNNER:-}" ]]; then
+      for WHEEL_PATH in dist/*.whl; do
+        WHEEL_NAME=$(basename "${WHEEL_PATH}")
+        aws s3 cp --acl public-read \
+          "${WHEEL_PATH}" \
+          "s3://gha-artifacts/${PYTORCH_WHEEL_S3_PATH}/${WHEEL_NAME}"
+      done
+    fi
+  fi
 
   # Grab the pinned audio and vision commits from PyTorch
   TORCHAUDIO_VERSION=$(cat .github/ci_commit_pins/audio.txt)
diff --git a/.github/workflows/_unittest.yml b/.github/workflows/_unittest.yml
index 6b08b6d1259..eb6c9c24257 100644
--- a/.github/workflows/_unittest.yml
+++ b/.github/workflows/_unittest.yml
@@ -49,4 +49,6 @@ jobs:
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
         set -eux
+        # This is needed to get the prebuilt PyTorch wheel from S3
+        ${CONDA_RUN} --no-capture-output pip install awscli==1.37.21
         .ci/scripts/unittest-macos.sh --build-tool "${{ inputs.build-tool }}" --build-mode "${{ inputs.build-mode }}" --editable "${{ inputs.editable }}"

From 3071d261be8c6d55bc8ffdd54f99c9c173999b0e Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Thu, 20 Mar 2025 21:43:52 -0700
Subject: [PATCH 2/9] No ACL

---
 .ci/scripts/utils.sh | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index d67faaa4b21..fa67bd577e0 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -77,9 +77,7 @@ install_pytorch_and_domains() {
     if command -v aws && [[ -z "${GITHUB_RUNNER:-}" ]]; then
       for WHEEL_PATH in dist/*.whl; do
         WHEEL_NAME=$(basename "${WHEEL_PATH}")
-        aws s3 cp --acl public-read \
-          "${WHEEL_PATH}" \
-          "s3://gha-artifacts/${PYTORCH_WHEEL_S3_PATH}/${WHEEL_NAME}"
+        aws s3 cp "${WHEEL_PATH}" "s3://gha-artifacts/${PYTORCH_WHEEL_S3_PATH}/${WHEEL_NAME}"
       done
     fi
   fi

From 947c39c94327dd958b7cd84ae8c0543dcd7fe008 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Fri, 21 Mar 2025 01:18:16 -0700
Subject: [PATCH 3/9] Implement reading from the cache

---
 .ci/scripts/utils.sh | 36 +++++++++++++++++++++++-------------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index fa67bd577e0..84c2cabcd01 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -55,19 +55,29 @@ install_pytorch_and_domains() {
   TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
   popd || return
 
-  OS_VERSION=$(uname)
-  PYTHON_VERSION=$(python -c 'import platform; print(platform.python_version())')
-  PYTORCH_WHEEL_S3_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${OS_VERSION}/${PYTHON_VERSION}"
-
-  if [[ "${FOUND_PYTORCH_WHEEL:-0}" == "0" ]]; then
-    # Found no such wheel, we will build it from source then
-    git clone https://github.com/pytorch/pytorch.git
-
-    # Fetch the target commit
-    pushd pytorch || return
-    git checkout "${TORCH_VERSION}"
-    git submodule update --init --recursive
+  git clone https://github.com/pytorch/pytorch.git
+
+  # Fetch the target commit
+  pushd pytorch || return
+  git checkout "${TORCH_VERSION}"
+  git submodule update --init --recursive
+
+  SYSTEM_NAME=$(uname)
+  PLATFORM=$(python -c 'import sysconfig; print(sysconfig.get_platform().replace("-", "_").replace(".", "_"))')
+  PYTHON_VERSION=$(python -c 'import platform; v = platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
+  TORCH_RELEASE=$(cat version.txt)
+  TORCH_SHORT_HASH=${TORCH_VERSION:0:7}
+  TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM}.whl"
+  TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}/"
+
+  # Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
+  # of the Docker image
+  if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
+    pip install "https://gha-artifacts.s3.us-east-1.amazonaws.com/${TORCH_WHEEL_PATH}/${TORCH_WHEEL_NAME}" || TORCH_WHEEL_NOT_FOUND=1
+  fi
 
+  # Found no such wheel, we will build it from source then
+  if [[ "${TORCH_WHEEL_NOT_FOUND:-0}" == "1" ]]; then
     export USE_DISTRIBUTED=1
     # Then build and install PyTorch
     python setup.py bdist_wheel
@@ -77,7 +87,7 @@ install_pytorch_and_domains() {
     if command -v aws && [[ -z "${GITHUB_RUNNER:-}" ]]; then
       for WHEEL_PATH in dist/*.whl; do
         WHEEL_NAME=$(basename "${WHEEL_PATH}")
-        aws s3 cp "${WHEEL_PATH}" "s3://gha-artifacts/${PYTORCH_WHEEL_S3_PATH}/${WHEEL_NAME}"
+        aws s3 cp "${WHEEL_PATH}" "s3://gha-artifacts/${TORCH_WHEEL_PATH}/${WHEEL_NAME}"
       done
     fi
   fi

From 9cbe26aa133110bb65777f19d152d1e97cf89020 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Fri, 21 Mar 2025 01:19:19 -0700
Subject: [PATCH 4/9] Fix typo

---
 .ci/scripts/utils.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index 84c2cabcd01..7d28b49b082 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -68,7 +68,7 @@ install_pytorch_and_domains() {
   TORCH_RELEASE=$(cat version.txt)
   TORCH_SHORT_HASH=${TORCH_VERSION:0:7}
   TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM}.whl"
-  TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}/"
+  TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}"
 
   # Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
   # of the Docker image

From 1576f322d0dbb9e755598255b2c9fc4843ff1ef2 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Fri, 21 Mar 2025 02:26:12 -0700
Subject: [PATCH 5/9] Another attempt

---
 .ci/scripts/utils.sh | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index 7d28b49b082..58d8fdf36ff 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -63,24 +63,24 @@ install_pytorch_and_domains() {
   git submodule update --init --recursive
 
   SYSTEM_NAME=$(uname)
-  PLATFORM=$(python -c 'import sysconfig; print(sysconfig.get_platform().replace("-", "_").replace(".", "_"))')
-  PYTHON_VERSION=$(python -c 'import platform; v = platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
+  # The platform version needs to match MACOSX_DEPLOYMENT_TARGET used to build the wheel
+  PLATFORM=$(python -c 'import sysconfig; platform=sysconfig.get_platform(); platform[1]="14_0"; print("_".join(platform))')
+  PYTHON_VERSION=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
   TORCH_RELEASE=$(cat version.txt)
   TORCH_SHORT_HASH=${TORCH_VERSION:0:7}
-  TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM}.whl"
   TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}"
+  TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM}.whl"
 
+  CACHE_TORCH_WHEEL="https://gha-artifacts.s3.us-east-1.amazonaws.com/${TORCH_WHEEL_PATH}/${TORCH_WHEEL_NAME}"
   # Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
   # of the Docker image
   if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
-    pip install "https://gha-artifacts.s3.us-east-1.amazonaws.com/${TORCH_WHEEL_PATH}/${TORCH_WHEEL_NAME}" || TORCH_WHEEL_NOT_FOUND=1
+    pip install "${CACHE_TORCH_WHEEL}" || TORCH_WHEEL_NOT_FOUND=1
   fi
 
   # Found no such wheel, we will build it from source then
   if [[ "${TORCH_WHEEL_NOT_FOUND:-0}" == "1" ]]; then
-    export USE_DISTRIBUTED=1
-    # Then build and install PyTorch
-    python setup.py bdist_wheel
+    USE_DISTRIBUTED=1 MACOSX_DEPLOYMENT_TARGET=14.0 python setup.py bdist_wheel
     pip install "$(echo dist/*.whl)"
 
     # Only AWS runners have access to S3
@@ -90,6 +90,8 @@ install_pytorch_and_domains() {
         aws s3 cp "${WHEEL_PATH}" "s3://gha-artifacts/${TORCH_WHEEL_PATH}/${WHEEL_NAME}"
       done
     fi
+  else
+    echo "Use cached wheel at ${CACHE_TORCH_WHEEL}"
   fi
 
   # Grab the pinned audio and vision commits from PyTorch

From fbd49491bc88b4942076359e401754e1635752cb Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Fri, 21 Mar 2025 02:34:11 -0700
Subject: [PATCH 6/9] Fix a bug

---
 .ci/scripts/utils.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index 58d8fdf36ff..fdc5ec33fef 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -64,7 +64,7 @@ install_pytorch_and_domains() {
 
   SYSTEM_NAME=$(uname)
   # The platform version needs to match MACOSX_DEPLOYMENT_TARGET used to build the wheel
-  PLATFORM=$(python -c 'import sysconfig; platform=sysconfig.get_platform(); platform[1]="14_0"; print("_".join(platform))')
+  PLATFORM=$(python -c 'import sysconfig; platform=sysconfig.get_platform().split("-"); platform[1]="14_0"; print("_".join(platform))')
   PYTHON_VERSION=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
   TORCH_RELEASE=$(cat version.txt)
   TORCH_SHORT_HASH=${TORCH_VERSION:0:7}

From bc6baf33b63083d5cb51c8cf77e00fd5e8e91a53 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Fri, 21 Mar 2025 14:37:10 -0700
Subject: [PATCH 7/9] Fix the platform version

---
 .ci/scripts/utils.sh        | 9 +++++----
 .github/workflows/trunk.yml | 2 +-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index fdc5ec33fef..25ef9f93e5b 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -63,13 +63,14 @@ install_pytorch_and_domains() {
   git submodule update --init --recursive
 
   SYSTEM_NAME=$(uname)
-  # The platform version needs to match MACOSX_DEPLOYMENT_TARGET used to build the wheel
-  PLATFORM=$(python -c 'import sysconfig; platform=sysconfig.get_platform().split("-"); platform[1]="14_0"; print("_".join(platform))')
+  if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
+    PLATFORM=$(python -c 'import sysconfig; import platform; v=platform.mac_ver()[0].split(".")[0]; platform=sysconfig.get_platform().split("-"); platform[1]=f"{v}_0"; print("_".join(platform))')
+  fi
   PYTHON_VERSION=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
   TORCH_RELEASE=$(cat version.txt)
   TORCH_SHORT_HASH=${TORCH_VERSION:0:7}
   TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}"
-  TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM}.whl"
+  TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM:-}.whl"
 
   CACHE_TORCH_WHEEL="https://gha-artifacts.s3.us-east-1.amazonaws.com/${TORCH_WHEEL_PATH}/${TORCH_WHEEL_NAME}"
   # Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
@@ -80,7 +81,7 @@ install_pytorch_and_domains() {
 
   # Found no such wheel, we will build it from source then
   if [[ "${TORCH_WHEEL_NOT_FOUND:-0}" == "1" ]]; then
-    USE_DISTRIBUTED=1 MACOSX_DEPLOYMENT_TARGET=14.0 python setup.py bdist_wheel
+    USE_DISTRIBUTED=1 python setup.py bdist_wheel
     pip install "$(echo dist/*.whl)"
 
     # Only AWS runners have access to S3
diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml
index 097a272d0fe..ecae932b74f 100644
--- a/.github/workflows/trunk.yml
+++ b/.github/workflows/trunk.yml
@@ -228,7 +228,7 @@ jobs:
     name: test-coreml-delegate
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
     with:
-      runner: macos-13-xlarge
+      runner: macos-latest-xlarge
       python-version: '3.11'
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}

From 3b1938f0d7d325d231694cf17375cd68ebd394e3 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Mon, 24 Mar 2025 12:06:56 -0700
Subject: [PATCH 8/9] Address review comments

---
 .ci/scripts/utils.sh | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index 25ef9f93e5b..c1af11bad58 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -60,39 +60,44 @@ install_pytorch_and_domains() {
   # Fetch the target commit
   pushd pytorch || return
   git checkout "${TORCH_VERSION}"
-  git submodule update --init --recursive
 
-  SYSTEM_NAME=$(uname)
-  if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
-    PLATFORM=$(python -c 'import sysconfig; import platform; v=platform.mac_ver()[0].split(".")[0]; platform=sysconfig.get_platform().split("-"); platform[1]=f"{v}_0"; print("_".join(platform))')
+  local system_name=$(uname)
+  if [[ "${system_name}" == "Darwin" ]]; then
+    local platform=$(python -c 'import sysconfig; import platform; v=platform.mac_ver()[0].split(".")[0]; platform=sysconfig.get_platform().split("-"); platform[1]=f"{v}_0"; print("_".join(platform))')
   fi
-  PYTHON_VERSION=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
-  TORCH_RELEASE=$(cat version.txt)
-  TORCH_SHORT_HASH=${TORCH_VERSION:0:7}
-  TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}"
-  TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM:-}.whl"
+  local python_version=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
+  local torch_release=$(cat version.txt)
+  local torch_short_hash=${TORCH_VERSION:0:7}
+  local torch_wheel_path="cached_artifacts/pytorch/executorch/pytorch_wheels/${system_name}/${python_version}"
+  local torch_wheel_name="torch-${torch_release}%2Bgit${torch_short_hash}-cp${python_version}-cp${python_version}-${platform:-}.whl"
 
-  CACHE_TORCH_WHEEL="https://gha-artifacts.s3.us-east-1.amazonaws.com/${TORCH_WHEEL_PATH}/${TORCH_WHEEL_NAME}"
+  local cached_torch_wheel="https://gha-artifacts.s3.us-east-1.amazonaws.com/${torch_wheel_path}/${torch_wheel_name}"
   # Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
   # of the Docker image
-  if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
-    pip install "${CACHE_TORCH_WHEEL}" || TORCH_WHEEL_NOT_FOUND=1
+  if [[ "${system_name}" == "Darwin" ]]; then
+    pip install "${cached_torch_wheel}" || TORCH_WHEEL_NOT_FOUND=1
+  else
+    TORCH_WHEEL_NOT_FOUND=1
   fi
 
   # Found no such wheel, we will build it from source then
   if [[ "${TORCH_WHEEL_NOT_FOUND:-0}" == "1" ]]; then
+    echo "No cached wheel found, continue with building PyTorch at ${TORCH_VERSION}"
+
+    git submodule update --init --recursive
     USE_DISTRIBUTED=1 python setup.py bdist_wheel
     pip install "$(echo dist/*.whl)"
 
     # Only AWS runners have access to S3
     if command -v aws && [[ -z "${GITHUB_RUNNER:-}" ]]; then
-      for WHEEL_PATH in dist/*.whl; do
-        WHEEL_NAME=$(basename "${WHEEL_PATH}")
-        aws s3 cp "${WHEEL_PATH}" "s3://gha-artifacts/${TORCH_WHEEL_PATH}/${WHEEL_NAME}"
+      for wheel_path in dist/*.whl; do
+        local wheel_name=$(basename "${wheel_path}")
+        echo "Caching ${wheel_name}"
+        aws s3 cp "${wheel_path}" "s3://gha-artifacts/${torch_wheel_path}/${wheel_name}"
       done
     fi
   else
-    echo "Use cached wheel at ${CACHE_TORCH_WHEEL}"
+    echo "Use cached wheel at ${cached_torch_wheel}"
   fi
 
   # Grab the pinned audio and vision commits from PyTorch

From 34fed00ca6bf97c4ce62983915cf65c0e6a912f5 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Mon, 24 Mar 2025 12:09:29 -0700
Subject: [PATCH 9/9] Another tweak

---
 .ci/scripts/utils.sh | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh
index c1af11bad58..677578ce3a4 100644
--- a/.ci/scripts/utils.sh
+++ b/.ci/scripts/utils.sh
@@ -74,14 +74,15 @@ install_pytorch_and_domains() {
   local cached_torch_wheel="https://gha-artifacts.s3.us-east-1.amazonaws.com/${torch_wheel_path}/${torch_wheel_name}"
   # Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
   # of the Docker image
+  local torch_wheel_not_found=0
   if [[ "${system_name}" == "Darwin" ]]; then
-    pip install "${cached_torch_wheel}" || TORCH_WHEEL_NOT_FOUND=1
+    pip install "${cached_torch_wheel}" || torch_wheel_not_found=1
   else
-    TORCH_WHEEL_NOT_FOUND=1
+    torch_wheel_not_found=1
   fi
 
   # Found no such wheel, we will build it from source then
-  if [[ "${TORCH_WHEEL_NOT_FOUND:-0}" == "1" ]]; then
+  if [[ "${torch_wheel_not_found}" == "1" ]]; then
     echo "No cached wheel found, continue with building PyTorch at ${TORCH_VERSION}"
 
     git submodule update --init --recursive