diff --git a/codegen/fix_includes.sh b/codegen/fix_includes.sh
index 9256fa2a9b8b..73a763a91cfb 100755
--- a/codegen/fix_includes.sh
+++ b/codegen/fix_includes.sh
@@ -6,7 +6,11 @@ sed -i 's#<torch_xla/csrc/tensor.h>#"torch_xla/csrc/tensor.h"#' $@
 # remove the runfiles-prefix used in codegen for pytorch
 # `torchgen` generates relative includes path and does not support customizing the root,
 # so we have to fix them up.
-sed -i 's#bazel-out/k8-[^/]*/bin/codegen/lazy_tensor_generator.runfiles/torch/##' $@
+if [[ $(uname -m) == "x86_64" ]]; then
+    sed -i 's#bazel-out/k8-[^/]*/bin/codegen/lazy_tensor_generator.runfiles/torch/##' $@
+elif [[ $(uname -m) == "aarch64" ]]; then
+    sed -i 's#bazel-out/aarch64-[^/]*/bin/codegen/lazy_tensor_generator.runfiles/torch/##' $@
+fi
 
 # use the generated files that are in the compilation unit
 sed -i 's#<bazel-out/\(.*\)>#"bazel-out/\1"#' $@
diff --git a/setup.py b/setup.py
index 9f82d66c32d3..3eae00e2796a 100644
--- a/setup.py
+++ b/setup.py
@@ -63,6 +63,10 @@
 
 import build_util
 
+import platform
+
+platform_machine = platform.machine()
+
 base_dir = os.path.dirname(os.path.abspath(__file__))
 
 USE_NIGHTLY = True  # whether to use nightly or stable libtpu and jax
@@ -77,13 +81,13 @@
   _libtpu_version += f".dev{_date}"
   _jax_version += f'.dev{_date}'
   _jaxlib_version += f'.dev{_date}'
-  _libtpu_wheel_name = f'libtpu-{_libtpu_version}.dev{_date}+nightly-py3-none-manylinux_2_31_x86_64'
+  _libtpu_wheel_name = f'libtpu-{_libtpu_version}.dev{_date}+nightly-py3-none-manylinux_2_31_{platform_machine}'
   _libtpu_storage_directory = 'libtpu-nightly-releases'
 else:
   # The postfix can be changed when the version is updated. Check
   # https://storage.googleapis.com/libtpu-wheels/index.html for correct
   # versioning.
-  _libtpu_wheel_name = f'libtpu-{_libtpu_version}-py3-none-manylinux_2_31_x86_64'
+  _libtpu_wheel_name = f'libtpu-{_libtpu_version}-py3-none-manylinux_2_31_{platform_machine}'
   _libtpu_storage_directory = 'libtpu-lts-releases'
 
 _libtpu_storage_path = f'https://storage.googleapis.com/{_libtpu_storage_directory}/wheels/libtpu/{_libtpu_wheel_name}.whl'