From 0c3b659e7709d760d52496076b26041c1e5706ea Mon Sep 17 00:00:00 2001
From: Robin Appelman <robin@icewind.nl>
Date: Mon, 24 Jun 2024 20:19:26 +0200
Subject: [PATCH] python3Packages.xformers: fix building with cudaSupport

---
 .../python-modules/xformers/default.nix           | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/xformers/default.nix b/pkgs/development/python-modules/xformers/default.nix
index b63270f6f4c5..261ff0cb89e8 100644
--- a/pkgs/development/python-modules/xformers/default.nix
+++ b/pkgs/development/python-modules/xformers/default.nix
@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   pythonOlder,
   fetchFromGitHub,
@@ -52,12 +53,14 @@ buildPythonPackage {
       # noqa: C801
       __version__ = "${version}"
       EOF
-    ''
-    + lib.optionalString cudaSupport ''
-      export CUDA_HOME=${cudaPackages.cuda_nvcc}
-      export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
     '';
 
+  env = lib.attrsets.optionalAttrs cudaSupport {
+    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+  };
+
+  stdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv;
+
   buildInputs = lib.optionals cudaSupport (
     with cudaPackages;
     [
@@ -71,7 +74,9 @@ buildPythonPackage {
     ]
   );
 
-  nativeBuildInputs = [ which ];
+  nativeBuildInputs = [ which ] ++ lib.optionals cudaSupport (with cudaPackages; [
+    cuda_nvcc
+  ]);
 
   propagatedBuildInputs = [
     numpy