python3Packages.onnxruntime: fix build with cudaSupport

2024-03-06 13:55:43 +01:00 · 2024-03-06 13:55:43 +01:00 · 05dedd387f
parent 7a4d4a55fc
commit 05dedd387f
2 changed files with 8 additions and 1 deletions
--- a/pkgs/development/libraries/onnxruntime/default.nix
+++ b/pkgs/development/libraries/onnxruntime/default.nix
@ -224,6 +224,7 @@ effectiveStdenv.mkDerivation rec {
  '';

  passthru = {
+    inherit cudaSupport cudaPackages; # for the python module
    protobuf = protobuf_21;
    tests = lib.optionalAttrs pythonSupport {
      python = python3Packages.onnxruntime;
--- a/pkgs/development/python-modules/onnxruntime/default.nix
+++ b/pkgs/development/python-modules/onnxruntime/default.nix
@ -53,7 +53,13 @@ buildPythonPackage {
    oneDNN
    re2
    onnxruntime.protobuf
-  ];
+  ] ++ lib.optionals onnxruntime.passthru.cudaSupport (with onnxruntime.passthru.cudaPackages; [
+    libcublas # libcublasLt.so.XX libcublas.so.XX
+    libcurand # libcurand.so.XX
+    libcufft # libcufft.so.XX
+    cudnn # libcudnn.soXX
+    cuda_cudart # libcudart.so.XX
+  ]);

  propagatedBuildInputs = [
    coloredlogs