@@ -1,4 +1,5 @@
11{ lib
2+, symlinkJoin
23, buildPythonPackage
34, fetchFromGitHub
45, ninja
@@ -10,9 +11,18 @@
1011, pillow
1112, pytorch
1213, pytest
14+, cudatoolkit
15+, cudnn
16+, cudaSupport ? pytorch.cudaSupport or false # by default uses the value from pytorch
1317}:
1418
15−buildPythonPackage rec {
19+let
20+ cudatoolkit_joined = symlinkJoin {
21+ name = "${cudatoolkit.name}-unsplit";
22+ paths = [ cudatoolkit.out cudatoolkit.lib ];
23+ };
24+ cudaArchStr = lib.optionalString cudaSupport lib.strings.concatStringsSep ";" pytorch.cudaArchList;
25+in buildPythonPackage rec {
1626 pname = "torchvision";
1727 version = "0.10.0";
1828
@@ -23,15 +33,22 @@ buildPythonPackage rec {
2333 sha256 = "13j04ij0jmi58nhav1p69xrm8dg7jisg23268i3n6lnms37n02kc";
2434 };
2535
26− nativeBuildInputs = [ libpng ninja which ];
36+ nativeBuildInputs = [ libpng ninja which ]
37+ ++ lib.optionals cudaSupport [ cudatoolkit_joined ];
2738
2839 TORCHVISION_INCLUDE = "${libjpeg_turbo.dev}/include/";
2940 TORCHVISION_LIBRARY = "${libjpeg_turbo}/lib/";
3041
31− buildInputs = [ libjpeg_turbo libpng ];
42+ buildInputs = [ libjpeg_turbo libpng ]
43+ ++ lib.optionals cudaSupport [ cudnn ];
3244
3345 propagatedBuildInputs = [ numpy pillow pytorch scipy ];
3446
47+ preBuild = lib.optionalString cudaSupport ''
48+ export TORCH_CUDA_ARCH_LIST="${cudaArchStr}"
49+ export FORCE_CUDA=1
50+ '';
51+
3552 # tries to download many datasets for tests
3653 doCheck = false;
3754
@@ -45,6 +62,7 @@ buildPythonPackage rec {
4562 description = "PyTorch vision library";
4663 homepage = "https://pytorch.org/";
4764 license = licenses.bsd3;
65+ platforms = with platforms; linux ++ lib.optionals (!cudaSupport) darwin;
4866 maintainers = with maintainers; [ ericsagnes ];
4967 };
5068}