about summary refs log tree commit diff
path: root/pkgs/development/python-modules/onnxruntime/default.nix
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/python-modules/onnxruntime/default.nix')
-rw-r--r--pkgs/development/python-modules/onnxruntime/default.nix8
1 files changed, 7 insertions, 1 deletions
diff --git a/pkgs/development/python-modules/onnxruntime/default.nix b/pkgs/development/python-modules/onnxruntime/default.nix
index 8fea2e6d9117e..0982a12c4572d 100644
--- 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