tensorflow: depend on cudatoolkit 8 and cudnn 5.1

This commit is contained in:
Jean-Philippe Bernardy 2017-02-27 15:23:23 +01:00
parent 7864782d71
commit f37b34fe7c
2 changed files with 8 additions and 6 deletions

View File

@ -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,7 +97,7 @@ 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
@ -105,7 +105,7 @@ buildPythonPackage rec {
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 ]

View File

@ -31256,6 +31256,8 @@ EOF
tensorflowWithCuda = callPackage ../development/python-modules/tensorflow {
cudaSupport = true;
cudatoolkit = pkgs.cudatoolkit8;
cudnn = pkgs.cudnn51_cudatoolkit80;
};
tflearn = buildPythonPackage rec {