Merge pull request #258563 from SomeoneSerge/fix/2023-09/torchaudio

python3Packages.torchaudio: fix the cuda build

pkgs/development/python-modules/torch/default.nix

@@ -438,11 +438,7 @@ in buildPythonPackage rec {
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-  } // lib.optionalAttrs cudaSupport {
-    # NOTE: supportedCudaCapabilities isn't computed unless cudaSupport is true, so we can't use
-    #   it in the passthru set above because a downstream package might try to access it even
-    #   when cudaSupport is false. Better to have it missing than null or an empty list by default.
-    cudaCapabilities = supportedCudaCapabilities;
+    cudaCapabilities = if cudaSupport then supportedCudaCapabilities else [ ];
   };
 
   meta = with lib; {

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