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