提交 54926a4b authored 作者: Frederic's avatar Frederic

set the cuda version when nvcc is not in the path but we use the cuda.root flag

上级 e51289c9
...@@ -60,14 +60,16 @@ nvcc_version = None ...@@ -60,14 +60,16 @@ nvcc_version = None
def is_nvcc_available(): def is_nvcc_available():
"""Return True iff the nvcc compiler is found.""" """Return True iff the nvcc compiler is found."""
try: def set_version():
p = subprocess.Popen(['nvcc', '--version'], stdout=subprocess.PIPE, p = subprocess.Popen([nvcc_path, '--version'], stdout=subprocess.PIPE,
stderr=subprocess.PIPE) stderr=subprocess.PIPE)
p.wait() p.wait()
s = p.stdout.readlines()[-1].split(',')[1].strip().split() s = p.stdout.readlines()[-1].split(',')[1].strip().split()
assert s[0] == 'release' assert s[0] == 'release'
global nvcc_version global nvcc_version
nvcc_version = s[1] nvcc_version = s[1]
try:
set_version()
return True return True
except Exception: except Exception:
#try to find nvcc into cuda.root #try to find nvcc into cuda.root
...@@ -75,6 +77,10 @@ def is_nvcc_available(): ...@@ -75,6 +77,10 @@ def is_nvcc_available():
if os.path.exists(p): if os.path.exists(p):
global nvcc_path global nvcc_path
nvcc_path = p nvcc_path = p
try:
set_version()
except Exception:
return False
return True return True
else: else:
return False return False
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论