diff options
| author | Alyssa Ross <hi@alyssa.is> | 2023-10-20 22:09:03 +0000 |
|---|---|---|
| committer | Alyssa Ross <hi@alyssa.is> | 2023-10-20 22:09:03 +0000 |
| commit | 50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e (patch) | |
| tree | f2556b911180125ccbb7ed0e78a54e92da89adce /nixpkgs/pkgs/development/python-modules/torchaudio/default.nix | |
| parent | 4c16d4548a98563c9d9ad76f4e5b2202864ccd54 (diff) | |
| parent | cfc75eec4603c06503ae750f88cf397e00796ea8 (diff) | |
| download | nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.tar nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.tar.gz nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.tar.bz2 nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.tar.lz nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.tar.xz nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.tar.zst nixlib-50c21d167f7114fa1dbd95e5c4fb30eeb1a2d02e.zip | |
Merge commit 'cfc75eec4603c06503ae750f88cf397e00796ea8'
Conflicts: nixpkgs/pkgs/build-support/rust/build-rust-package/default.nix
Diffstat (limited to 'nixpkgs/pkgs/development/python-modules/torchaudio/default.nix')
| -rw-r--r-- | nixpkgs/pkgs/development/python-modules/torchaudio/default.nix | 19 |
1 files changed, 16 insertions, 3 deletions
diff --git a/nixpkgs/pkgs/development/python-modules/torchaudio/default.nix b/nixpkgs/pkgs/development/python-modules/torchaudio/default.nix index 3bd8003890d7..0b38925e0a2b 100644 --- a/nixpkgs/pkgs/development/python-modules/torchaudio/default.nix +++ b/nixpkgs/pkgs/development/python-modules/torchaudio/default.nix @@ -6,7 +6,7 @@ , ninja , pybind11 , torch -, cudaSupport ? false +, cudaSupport ? torch.cudaSupport , cudaPackages }: @@ -27,17 +27,30 @@ buildPythonPackage rec { --replace "_fetch_archives(_parse_sources())" "pass" ''; + env = { + TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}"; + }; + nativeBuildInputs = [ cmake pkg-config ninja ] ++ lib.optionals cudaSupport [ - cudaPackages.cudatoolkit + cudaPackages.cuda_nvcc ]; buildInputs = [ pybind11 ] ++ lib.optionals cudaSupport [ - cudaPackages.cudnn + cudaPackages.libcurand.dev + cudaPackages.libcurand.lib + cudaPackages.cuda_cudart # cuda_runtime.h and libraries + cudaPackages.cuda_cccl.dev # <thrust/*> + cudaPackages.cuda_nvtx.dev + cudaPackages.cuda_nvtx.lib # -llibNVToolsExt + cudaPackages.libcublas.dev + cudaPackages.libcublas.lib + cudaPackages.libcufft.dev + cudaPackages.libcufft.lib ]; propagatedBuildInputs = [ torch |