pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, cmake
5, pkg-config
6, ninja
7, pybind11
8, torch
9, cudaSupport ? torch.cudaSupport
10, cudaPackages
11}:
12
13buildPythonPackage rec {
14 pname = "torchaudio";
15 version = "2.0.2";
16
17 src = fetchFromGitHub {
18 owner = "pytorch";
19 repo = "audio";
20 rev = "v${version}";
21 hash = "sha256-9lB4gLXq0nXHT1+DNOlbJQqNndt2I6kVoNwhMO/2qlE=";
22 };
23
24 postPatch = ''
25 substituteInPlace setup.py \
26 --replace 'print(" --- Initializing submodules")' "return" \
27 --replace "_fetch_archives(_parse_sources())" "pass"
28 '';
29
30 env = {
31 TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
32 };
33
34 nativeBuildInputs = [
35 cmake
36 pkg-config
37 ninja
38 ] ++ lib.optionals cudaSupport [
39 cudaPackages.cuda_nvcc
40 ];
41 buildInputs = [
42 pybind11
43 ] ++ lib.optionals cudaSupport [
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
54 ];
55 propagatedBuildInputs = [
56 torch
57 ];
58
59 BUILD_SOX=0;
60 BUILD_KALDI=0;
61 BUILD_RNNT=0;
62 BUILD_CTC_DECODER=0;
63
64 dontUseCmakeConfigure = true;
65
66 doCheck = false; # requires sox backend
67
68 meta = with lib; {
69 description = "PyTorch audio library";
70 homepage = "https://pytorch.org/";
71 changelog = "https://github.com/pytorch/audio/releases/tag/v${version}";
72 license = licenses.bsd2;
73 platforms = platforms.unix;
74 maintainers = with maintainers; [ junjihashimoto ];
75 };
76}