about summary refs log tree commit diff
path: root/pkgs/development/python-modules/torch/prefetch.sh
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/python-modules/torch/prefetch.sh')
-rwxr-xr-xpkgs/development/python-modules/torch/prefetch.sh11
1 files changed, 6 insertions, 5 deletions
diff --git a/pkgs/development/python-modules/torch/prefetch.sh b/pkgs/development/python-modules/torch/prefetch.sh
index 31aae438dc56d..78151ae787a52 100755
--- a/pkgs/development/python-modules/torch/prefetch.sh
+++ b/pkgs/development/python-modules/torch/prefetch.sh
@@ -5,14 +5,15 @@ set -eou pipefail
 
 version=$1
 
-linux_bucket="https://download.pytorch.org/whl/cu113"
+linux_cuda_version="cu116"
+linux_bucket="https://download.pytorch.org/whl/${linux_cuda_version}"
 darwin_bucket="https://download.pytorch.org/whl/cpu"
 
 url_and_key_list=(
-  "x86_64-linux-37 $linux_bucket/torch-${version}%2Bcu113-cp37-cp37m-linux_x86_64.whl torch-${version}-cp37-cp37m-linux_x86_64.whl"
-  "x86_64-linux-38 $linux_bucket/torch-${version}%2Bcu113-cp38-cp38-linux_x86_64.whl torch-${version}-cp38-cp38-linux_x86_64.whl"
-  "x86_64-linux-39 $linux_bucket/torch-${version}%2Bcu113-cp39-cp39-linux_x86_64.whl torch-${version}-cp39-cp39-linux_x86_64.whl"
-  "x86_64-linux-310 $linux_bucket/torch-${version}%2Bcu113-cp310-cp310-linux_x86_64.whl torch-${version}-cp310-cp310-linux_x86_64.whl"
+  "x86_64-linux-37 $linux_bucket/torch-${version}%2B${linux_cuda_version}-cp37-cp37m-linux_x86_64.whl torch-${version}-cp37-cp37m-linux_x86_64.whl"
+  "x86_64-linux-38 $linux_bucket/torch-${version}%2B${linux_cuda_version}-cp38-cp38-linux_x86_64.whl torch-${version}-cp38-cp38-linux_x86_64.whl"
+  "x86_64-linux-39 $linux_bucket/torch-${version}%2B${linux_cuda_version}-cp39-cp39-linux_x86_64.whl torch-${version}-cp39-cp39-linux_x86_64.whl"
+  "x86_64-linux-310 $linux_bucket/torch-${version}%2B${linux_cuda_version}-cp310-cp310-linux_x86_64.whl torch-${version}-cp310-cp310-linux_x86_64.whl"
   "x86_64-darwin-37 $darwin_bucket/torch-${version}-cp37-none-macosx_10_9_x86_64.whl torch-${version}-cp37-none-macosx_10_9_x86_64.whl"
   "x86_64-darwin-38 $darwin_bucket/torch-${version}-cp38-none-macosx_10_9_x86_64.whl torch-${version}-cp38-none-macosx_10_9_x86_64.whl"
   "x86_64-darwin-39 $darwin_bucket/torch-${version}-cp39-none-macosx_10_9_x86_64.whl torch-${version}-cp39-none-macosx_10_9_x86_64.whl"