python37Packages.cupy: Add cutensor optional dep

NixOS · Mar 17, 2021 · 13b5995 · 13b5995
1 parent a680b02
commit 13b5995
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/pkgs/development/python-modules/cupy/default.nix b/pkgs/development/python-modules/cupy/default.nix
@@ -1,7 +1,7 @@
 { lib, buildPythonPackage
 , fetchPypi, isPy3k, linuxPackages
 , fastrlock, numpy, six, wheel, pytest, mock, setuptools
-, cudatoolkit, cudnn, nccl
+, cudatoolkit, cudnn, cutensor, nccl
 }:
 
 buildPythonPackage rec {
@@ -26,6 +26,7 @@ buildPythonPackage rec {
   propagatedBuildInputs = [
     cudatoolkit
     cudnn
+    cutensor
     linuxPackages.nvidia_x11
     nccl
     fastrlock

diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
@@ -1589,6 +1589,7 @@ in {
     cudatoolkit = pkgs.cudatoolkit_10_0;
     cudnn = pkgs.cudnn_cudatoolkit_10_0;
     nccl = pkgs.nccl_cudatoolkit_10;
+    cutensor = pkgs.cutensor_cudatoolkit_10;
   };
 
   curio = callPackage ../development/python-modules/curio { };