diff options
Diffstat (limited to 'pkgs/development/python-modules/onnxruntime/default.nix')
-rw-r--r-- | pkgs/development/python-modules/onnxruntime/default.nix | 8 |
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 |