about summary refs log tree commit diff
path: root/pkgs/development/libraries/science/math/libtorch/prefetch.sh
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/libraries/science/math/libtorch/prefetch.sh')
-rwxr-xr-xpkgs/development/libraries/science/math/libtorch/prefetch.sh5
1 files changed, 3 insertions, 2 deletions
diff --git a/pkgs/development/libraries/science/math/libtorch/prefetch.sh b/pkgs/development/libraries/science/math/libtorch/prefetch.sh
index 5c6d60ae8b208..2682fe6779768 100755
--- a/pkgs/development/libraries/science/math/libtorch/prefetch.sh
+++ b/pkgs/development/libraries/science/math/libtorch/prefetch.sh
@@ -6,10 +6,11 @@ set -eou pipefail
 version=$1
 
 bucket="https://download.pytorch.org/libtorch"
-CUDA_VERSION=cu116
+CUDA_VERSION=cu121
 
 url_and_key_list=(
-  "x86_64-darwin-cpu $bucket/cpu/libtorch-macos-${version}.zip libtorch-macos-${version}.zip"
+  "aarch64-darwin-cpu $bucket/cpu/libtorch-macos-arm64-${version}.zip libtorch-macos-arm64-${version}.zip"
+  "x86_64-darwin-cpu $bucket/cpu/libtorch-macos-x86_64-${version}.zip libtorch-macos-x86_64-${version}.zip"
   "x86_64-linux-cpu $bucket/cpu/libtorch-cxx11-abi-shared-with-deps-${version}%2Bcpu.zip libtorch-cxx11-abi-shared-with-deps-${version}-cpu.zip"
   "x86_64-linux-cuda $bucket/${CUDA_VERSION}/libtorch-cxx11-abi-shared-with-deps-${version}%2B${CUDA_VERSION}.zip libtorch-cxx11-abi-shared-with-deps-${version}-${CUDA_VERSION}.zip"
 )