1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-14 00:06:58 +08:00
This commit is contained in:
Jactus
2020-11-26 13:11:56 +08:00
parent 37cc51465c
commit 45192413f9
2 changed files with 5 additions and 4 deletions

View File

@@ -164,13 +164,14 @@ class ExtendedEnvBuilder(venv.EnvBuilder):
# function to check cuda version on the machine, this case is for the model TFT # function to check cuda version on the machine, this case is for the model TFT
def check_cuda(folders): def check_cuda(folders):
path = "/usr/local/cuda/version.txt" path = "/usr/local/cuda/version.txt" # TODO: FIX ME, this will not work on other os systems.
exclude_tft = True exclude_tft = True
if os.path.exists(path): if os.path.exists(path):
with open(path, "w") as f: with open(path, "r") as f:
if "10.1" in str(f.read()) or "10.0" in str(f.read()): if "10.1" in str(f.read()) or "10.0" in str(f.read()):
exclude_tft = False exclude_tft = False
if exclude_tft and "TFT" in folders: if exclude_tft and "TFT" in folders:
sys.stderr.write("Compatible CUDA version not found! Removing TFT from the workflow...\n")
del folders["TFT"] del folders["TFT"]
return folders return folders

View File

@@ -27,9 +27,9 @@ def sys_config(config, config_path):
Parameters Parameters
---------- ----------
config : dict config : dict
configuration of the workflow configuration of the workflow.
config_path : str config_path : str
configuration of the path configuration of the path.
""" """
sys_config = config.get("sys", {}) sys_config = config.get("sys", {})