diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 43d6a4f..232233f 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -6,9 +6,17 @@ find_package(Python 3.10 REQUIRED COMPONENTS Interpreter Development.Module) # Download LibTorch include(FetchContent) +if (AARCH64 OR CMAKE_SYSTEM_PROCESSOR MATCHES "(arm)|(ARM)|(aarch64)") + message(STATUS "Targeting ARM architecture (AARCH64 override or processor match)") + set(LIBTORCH_URL "https://github.com/PASSIONLab/libtorch_aarch64/raw/refs/heads/main/pkg-libtorch.tar.zst") +else() + message(STATUS "Did not match arm/aarch64, performing x86 install") + set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.10.0%2Bcpu.zip") +endif() + FetchContent_Declare( libtorch - URL "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.10.0%2Bcpu.zip" + URL "${LIBTORCH_URL}" ) message(STATUS "Downloading LibTorch...") diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 4085681..e58669c 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "openequivariance" -version = "0.6.5" +version = "0.6.6" authors = [ { name="Austin Glover" }, { name="Vivek Bharadwaj" }, @@ -75,3 +75,6 @@ addopts = [ [tool.ruff] lint.ignore = ["E741"] + +[tool.scikit-build.cmake.define] +AARCH64 = {env="AARCH64", default="0"} diff --git a/openequivariance_extjax/pyproject.toml b/openequivariance_extjax/pyproject.toml index ba313d2..af35dde 100644 --- a/openequivariance_extjax/pyproject.toml +++ b/openequivariance_extjax/pyproject.toml @@ -8,7 +8,7 @@ build-backend = "scikit_build_core.build" [project] name = "openequivariance_extjax" -version = "0.6.5" +version = "0.6.6" authors = [ { name="Austin Glover" },