tensorflow: depend on cudatoolkit 8 and cudnn 5.1 (#24686)

(cherry-picked from commit f37b34f)
NixOS · Apr 7, 2017 · ab5d7e0 · ab5d7e0
1 parent 536d3ca
commit ab5d7e0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix
@@ -3,8 +3,8 @@
 , buildPythonPackage
 , isPy35, isPy27
 , cudaSupport ? false
-, cudatoolkit75 ? null
-, cudnn5_cudatoolkit75 ? null
+, cudatoolkit ? null
+, cudnn ? null
 , gcc49 ? null
 , linuxPackages ? null
 , numpy
@@ -16,8 +16,8 @@
 , zlib
 }:
 
-assert cudaSupport -> cudatoolkit75 != null
-                   && cudnn5_cudatoolkit75 != null
+assert cudaSupport -> cudatoolkit != null
+                   && cudnn != null
                    && gcc49 != null
                    && linuxPackages != null;
 
@@ -97,15 +97,15 @@ buildPythonPackage rec {
 
   propagatedBuildInputs = with stdenv.lib;
     [ numpy six protobuf3_2 swig mock ]
-    ++ optionals cudaSupport [ cudatoolkit75 cudnn5_cudatoolkit75 gcc49 ];
+    ++ optionals cudaSupport [ cudatoolkit cudnn gcc49 ];
 
   # Note that we need to run *after* the fixup phase because the
   # libraries are loaded at runtime. If we run in preFixup then
   # patchelf --shrink-rpath will remove the cuda libraries.
   postFixup = let
     rpath = stdenv.lib.makeLibraryPath
       (if cudaSupport then
-        [ gcc49.cc.lib zlib cudatoolkit75 cudnn5_cudatoolkit75
+        [ gcc49.cc.lib zlib cudatoolkit cudnn
           linuxPackages.nvidia_x11 ]
       else
         [ gcc.cc.lib zlib ]

diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
@@ -31219,6 +31219,8 @@ EOF
 
   tensorflowWithCuda = callPackage ../development/python-modules/tensorflow {
     cudaSupport = true;
+    cudatoolkit = pkgs.cudatoolkit8;
+    cudnn = pkgs.cudnn51_cudatoolkit80;
   };
 
   tflearn = buildPythonPackage rec {