From 3ed81c19e97d5ad4f40cd0973cb7c8562976c634 Mon Sep 17 00:00:00 2001 From: Benjamin Berman Date: Fri, 4 Aug 2023 16:26:13 -0700 Subject: [PATCH] select bog standard for linux arm64 since nightlies aren't built for it --- setup.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/setup.py b/setup.py index 09464551e..bce235833 100644 --- a/setup.py +++ b/setup.py @@ -41,6 +41,12 @@ This includes macOS MPS support. """ cpu_torch_index_nightlies = "https://download.pytorch.org/whl/nightly/cpu" +""" +The package index to torch built against CPU features. +Non-nightlies are selected when building Linux on arm64. +""" +cpu_torch_index = "https://download.pytorch.org/whl/cpu" + # xformers not required for new torch """ @@ -98,6 +104,13 @@ def _is_amd() -> bool: return False +def _is_linux_arm64(): + os_name = platform.system() + architecture = platform.machine() + + return os_name == 'Linux' and architecture == 'aarch64' + + def dependencies() -> [str]: _dependencies = open(os.path.join(os.path.dirname(__file__), "requirements.txt")).readlines() # todo: also add all plugin dependencies @@ -113,6 +126,8 @@ def dependencies() -> [str]: elif _is_amd(): index_urls += [amd_torch_index] gpu_accelerated = True + elif _is_linux_arm64(): + index_urls += [cpu_torch_index] else: index_urls += [cpu_torch_index_nightlies]