python3Packages.mxnet: pass explicit cuda gencodes
This commit is contained in:
parent
e5e38773a5
commit
226fabb1a6
|
@ -2,6 +2,7 @@
|
|||
, opencv3, gtest, blas, gomp, llvmPackages, perl
|
||||
, cudaSupport ? config.cudaSupport or false, cudatoolkit, nvidia_x11
|
||||
, cudnnSupport ? cudaSupport, cudnn
|
||||
, cudaCapabilities ? [ "3.7" "5.0" "6.0" "7.0" "7.5" "8.0" "8.6" ]
|
||||
}:
|
||||
|
||||
assert cudnnSupport -> cudaSupport;
|
||||
|
@ -44,6 +45,7 @@ stdenv.mkDerivation rec {
|
|||
"-DUSE_OLDCMAKECUDA=ON" # see https://github.com/apache/incubator-mxnet/issues/10743
|
||||
"-DCUDA_ARCH_NAME=All"
|
||||
"-DCUDA_HOST_COMPILER=${cudatoolkit.cc}/bin/cc"
|
||||
"-DMXNET_CUDA_ARCH=${lib.concatStringsSep ";" cudaCapabilities}"
|
||||
] else [ "-DUSE_CUDA=OFF" ])
|
||||
++ lib.optional (!cudnnSupport) "-DUSE_CUDNN=OFF";
|
||||
|
||||
|
|
Loading…
Reference in a new issue