nordic-dev.net/nixpkgs
nordic-dev.net/nixpkgs
2
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2024-12-16 10:43:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

python3Packages.torch: update CUDA capabilities for v2.0.1 release

Browse Source
This commit is contained in:
Connor Baker 2023-05-25 18:02:05 -04:00
parent 178b88e3ae
commit a52e068d86
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
pkgs/development/python-modules/torch/default.nix
View File

@ -65,10 +65,10 @@ assert !cudaSupport || magma.cudaPackages.cudatoolkit == cudatoolkit;
let
setBool = v: if v then "1" else "0";
# https://github.com/pytorch/pytorch/blob/v1.13.1/torch/utils/cpp_extension.py#L1751
# https://github.com/pytorch/pytorch/blob/v2.0.1/torch/utils/cpp_extension.py#L1744
supportedTorchCudaCapabilities =
let
real = ["3.5" "3.7" "5.0" "5.2" "5.3" "6.0" "6.1" "6.2" "7.0" "7.2" "7.5" "8.0" "8.6"];
real = ["3.5" "3.7" "5.0" "5.2" "5.3" "6.0" "6.1" "6.2" "7.0" "7.2" "7.5" "8.0" "8.6" "8.9" "9.0"];
ptx = lists.map (x: "${x}+PTX") real;
in
real ++ ptx;