pyrox.dev
lol
python3Packages.torchaudio: fix the cuda build
+16
-3
pkgs/development/python-modules/torchaudio/default.nix
+16
-3
pkgs/development/python-modules/torchaudio/default.nix
···
6
6
, ninja
7
7
, pybind11
8
8
, torch
9
-
, cudaSupport ? false
9
+
, cudaSupport ? torch.cudaSupport
10
10
, cudaPackages
11
11
}:
12
12
···
27
27
--replace "_fetch_archives(_parse_sources())" "pass"
28
28
'';
29
29
30
+
env = {
31
+
TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
32
+
};
33
+
30
34
nativeBuildInputs = [
31
35
cmake
32
36
pkg-config
33
37
ninja
34
38
] ++ lib.optionals cudaSupport [
35
-
cudaPackages.cudatoolkit
39
+
cudaPackages.cuda_nvcc
36
40
];
37
41
buildInputs = [
38
42
pybind11
39
43
] ++ lib.optionals cudaSupport [
40
-
cudaPackages.cudnn
44
+
cudaPackages.libcurand.dev
45
+
cudaPackages.libcurand.lib
46
+
cudaPackages.cuda_cudart # cuda_runtime.h and libraries
47
+
cudaPackages.cuda_cccl.dev # <thrust/*>
48
+
cudaPackages.cuda_nvtx.dev
49
+
cudaPackages.cuda_nvtx.lib # -llibNVToolsExt
50
+
cudaPackages.libcublas.dev
51
+
cudaPackages.libcublas.lib
52
+
cudaPackages.libcufft.dev
53
+
cudaPackages.libcufft.lib
41
54
];
42
55
propagatedBuildInputs = [
43
56
torch