PyTorch: specify CUDA root directory (#40855)
This commit is contained in:
parent
dd2b436b5a
commit
b8bb8a70ce
1 changed files with 2 additions and 3 deletions
|
@ -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
|
||||||
)
|
)
|
||||||
|
|
Loading…
Reference in a new issue