PyTorch: specify CUDA root directory (#40855)

This commit is contained in:
Adam J. Stewart 2023-11-09 16:25:54 -06:00 committed by GitHub
parent dd2b436b5a
commit b8bb8a70ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -491,9 +491,8 @@ def enable_or_disable(variant, keyword="USE", var=None, newer=False):
enable_or_disable("cuda") enable_or_disable("cuda")
if "+cuda" in self.spec: if "+cuda" in self.spec:
# cmake/public/cuda.cmake env.set("CUDA_HOME", self.spec["cuda"].prefix) # Linux/macOS
# cmake/Modules_CUDA_fix/upstream/FindCUDA.cmake env.set("CUDA_PATH", self.spec["cuda"].prefix) # Windows
env.unset("CUDA_ROOT")
torch_cuda_arch = ";".join( torch_cuda_arch = ";".join(
"{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value "{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value
) )