From 55e0d9a93dd8946fb5bb0008cd00c90f161cac12 Mon Sep 17 00:00:00 2001
From: Connor Baker <connorbaker01@gmail.com>
Date: Sat, 11 Mar 2023 00:43:21 -0500
Subject: [PATCH 1/2] torch: passthru.cudaCapabilities is available only when
 cudaSupport is true

Additionally, move away from gpuTargetString to the more flexible cudaCapabilities.
---
 pkgs/development/python-modules/torch/default.nix       | 8 ++++++--
 pkgs/development/python-modules/torchvision/default.nix | 4 ++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index b00146b573fe..27e014201d65 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -360,10 +360,14 @@ in buildPythonPackage rec {
   requiredSystemFeatures = [ "big-parallel" ];
 
   passthru = {
-    inherit cudaSupport cudaPackages gpuTargetString;
-    cudaCapabilities = supportedCudaCapabilities;
+    inherit cudaSupport cudaPackages;
     # At least for 1.10.2 `torch.fft` is unavailable unless BLAS provider is MKL. This attribute allows for easy detection of its availability.
     blasProvider = blas.provider;
+  } // 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;
   };
 
   meta = with lib; {
diff --git a/pkgs/development/python-modules/torchvision/default.nix b/pkgs/development/python-modules/torchvision/default.nix
index 8201d9c72b73..eef6acd8cc23 100644
--- a/pkgs/development/python-modules/torchvision/default.nix
+++ b/pkgs/development/python-modules/torchvision/default.nix
@@ -15,7 +15,7 @@
 }:
 
 let
-  inherit (torch) cudaPackages gpuTargetString;
+  inherit (torch) cudaPackages cudaCapabilities;
   inherit (cudaPackages) cudatoolkit cudaFlags cudaVersion;
 
   # NOTE: torchvision doesn't use cudnn; torch does!
@@ -68,7 +68,7 @@ buildPythonPackage {
   + lib.optionalString cudaSupport ''
     export CC=${cudatoolkit.cc}/bin/cc
     export CXX=${cudatoolkit.cc}/bin/c++
-    export TORCH_CUDA_ARCH_LIST="${gpuTargetString}"
+    export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
     export FORCE_CUDA=1
   '';
 

From 504d7531bac0560cf1a8c6aa995e61b234aa9a40 Mon Sep 17 00:00:00 2001
From: Connor Baker <connorbaker01@gmail.com>
Date: Sat, 11 Mar 2023 00:55:50 -0500
Subject: [PATCH 2/2] torchvision: remove cudaSupport argument; must use
 whatever torch used

---
 pkgs/development/python-modules/torchvision/default.nix | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/torchvision/default.nix b/pkgs/development/python-modules/torchvision/default.nix
index eef6acd8cc23..9539b1bb0bfe 100644
--- a/pkgs/development/python-modules/torchvision/default.nix
+++ b/pkgs/development/python-modules/torchvision/default.nix
@@ -1,5 +1,4 @@
 { buildPythonPackage
-, cudaSupport ? torch.cudaSupport or false # by default uses the value from torch
 , fetchFromGitHub
 , lib
 , libjpeg_turbo
@@ -15,7 +14,7 @@
 }:
 
 let
-  inherit (torch) cudaPackages cudaCapabilities;
+  inherit (torch) cudaCapabilities cudaPackages cudaSupport;
   inherit (cudaPackages) cudatoolkit cudaFlags cudaVersion;
 
   # NOTE: torchvision doesn't use cudnn; torch does!