From a6cd76d7e992db61629b7ee23766bc9fa886ca62 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 2 Oct 2023 13:25:53 +0300
Subject: [PATCH 1/2] python3Packages.torchaudio: fix the cuda build

---
 .../python-modules/torchaudio/default.nix     | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/torchaudio/default.nix b/pkgs/development/python-modules/torchaudio/default.nix
index 3bd8003890d7..0b38925e0a2b 100644
--- a/pkgs/development/python-modules/torchaudio/default.nix
+++ b/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

From 7e0b4a80205a64750873e53a5eea0e86d757c5f2 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 2 Oct 2023 14:48:56 +0300
Subject: [PATCH 2/2] fixup! python3Packages.torchaudio: fix the cuda build

---
 pkgs/development/python-modules/torch/default.nix | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index f9f6e377b139..fa74bbebfd89 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -436,11 +436,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; {