nordic-dev.net/nixpkgs
nordic-dev.net/nixpkgs
2
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2025-02-15 16:45:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

python3Packages.xformers: fix building with cudaSupport

Browse Source
This commit is contained in:
Robin Appelman 2024-06-24 20:19:26 +02:00
parent 0645ee4c27
commit 0c3b659e77
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
pkgs/development/python-modules/xformers/default.nix
View File

@ -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