summary refs log tree commit diff
path: root/pkgs/development/python-modules/tensorflow/bin.nix
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/python-modules/tensorflow/bin.nix')
-rw-r--r--pkgs/development/python-modules/tensorflow/bin.nix10
1 files changed, 9 insertions, 1 deletions
diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
index e46d0298953..db2308a0be8 100644
--- a/pkgs/development/python-modules/tensorflow/bin.nix
+++ b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -26,6 +26,7 @@
 , symlinkJoin
 , keras-applications
 , keras-preprocessing
+, addOpenGLRunpath
 }:
 
 # We keep this binary build for two reasons:
@@ -76,6 +77,8 @@ in buildPythonPackage {
   ] ++ lib.optional (!isPy3k) mock
     ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
 
+  nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
+
   # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
   # and the propageted input tensorflow-tensorboard which causes environment collisions.
   # another possibility would be to have tensorboard only in the buildInputs
@@ -94,7 +97,12 @@ in buildPythonPackage {
   lib.optionalString stdenv.isLinux ''
     rrPath="$out/${python.sitePackages}/tensorflow/:$out/${python.sitePackages}/tensorflow/contrib/tensor_forest/:${rpath}"
     internalLibPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so"
-    find $out \( -name '*.so' -or -name '*.so.*' \) -exec patchelf --set-rpath "$rrPath" {} \;
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      patchelf --set-rpath "$rrPath" "$lib"
+      ${lib.optionalString cudaSupport ''
+        addOpenGLRunpath "$lib"
+      ''}
+    done
   '';