python37Packages.cupy: Add cutensor optional dep

This commit is contained in:
John Ericson 2021-03-15 23:09:15 +00:00
parent a680b02816
commit 13b5995f06
2 changed files with 3 additions and 1 deletions

View File

@ -1,7 +1,7 @@
{ lib, buildPythonPackage { lib, buildPythonPackage
, fetchPypi, isPy3k, linuxPackages , fetchPypi, isPy3k, linuxPackages
, fastrlock, numpy, six, wheel, pytest, mock, setuptools , fastrlock, numpy, six, wheel, pytest, mock, setuptools
, cudatoolkit, cudnn, nccl , cudatoolkit, cudnn, cutensor, nccl
}: }:
buildPythonPackage rec { buildPythonPackage rec {
@ -26,6 +26,7 @@ buildPythonPackage rec {
propagatedBuildInputs = [ propagatedBuildInputs = [
cudatoolkit cudatoolkit
cudnn cudnn
cutensor
linuxPackages.nvidia_x11 linuxPackages.nvidia_x11
nccl nccl
fastrlock fastrlock

View File

@ -1589,6 +1589,7 @@ in {
cudatoolkit = pkgs.cudatoolkit_10_0; cudatoolkit = pkgs.cudatoolkit_10_0;
cudnn = pkgs.cudnn_cudatoolkit_10_0; cudnn = pkgs.cudnn_cudatoolkit_10_0;
nccl = pkgs.nccl_cudatoolkit_10; nccl = pkgs.nccl_cudatoolkit_10;
cutensor = pkgs.cutensor_cudatoolkit_10;
}; };
curio = callPackage ../development/python-modules/curio { }; curio = callPackage ../development/python-modules/curio { };