Merge pull request #221950 from ConnorBaker/fix/221898-torchvision-backendStdenv

python3Packages.torchvision: switch to backendStdenv.cc from cudatoolkit.cc
This commit is contained in:
Samuel Ainsworth 2023-03-20 17:35:30 -04:00 committed by GitHub
commit 862746a32f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -15,7 +15,7 @@
let
inherit (torch) cudaCapabilities cudaPackages cudaSupport;
inherit (cudaPackages) cudatoolkit cudaFlags cudaVersion;
inherit (cudaPackages) backendStdenv cudaVersion;
# NOTE: torchvision doesn't use cudnn; torch does!
# For this reason it is not included.
@ -65,8 +65,8 @@ buildPythonPackage {
# NOTE: We essentially override the compilers provided by stdenv because we don't have a hook
# for cudaPackages to swap in compilers supported by NVCC.
+ lib.optionalString cudaSupport ''
export CC=${cudatoolkit.cc}/bin/cc
export CXX=${cudatoolkit.cc}/bin/c++
export CC=${backendStdenv.cc}/bin/cc
export CXX=${backendStdenv.cc}/bin/c++
export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
export FORCE_CUDA=1
'';