mirror of
https://github.com/NixOS/nixpkgs.git
synced 2024-11-16 23:03:40 +01:00
tensorflow: depend on cudatoolkit 8 and cudnn 5.1
This commit is contained in:
parent
7864782d71
commit
f37b34fe7c
2 changed files with 8 additions and 6 deletions
|
@ -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 ]
|
||||
|
|
|
@ -31256,6 +31256,8 @@ EOF
|
|||
|
||||
tensorflowWithCuda = callPackage ../development/python-modules/tensorflow {
|
||||
cudaSupport = true;
|
||||
cudatoolkit = pkgs.cudatoolkit8;
|
||||
cudnn = pkgs.cudnn51_cudatoolkit80;
|
||||
};
|
||||
|
||||
tflearn = buildPythonPackage rec {
|
||||
|
|
Loading…
Reference in a new issue