only change cuda and openvino torch command

This commit is contained in:
Yoink3000
2024-10-22 04:39:50 +08:00
parent 09a4320ab8
commit 26ae378061
+4 -4
View File
@@ -489,7 +489,7 @@ def install_cuda():
log.info('CUDA: nVidia toolkit detected')
install('onnxruntime-gpu', 'onnxruntime-gpu', ignore=True, quiet=True)
# return os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/cu124')
return os.environ.get('TORCH_COMMAND', 'torch==2.4.1+cu124 torchvision==0.19.1 --index-url https://download.pytorch.org/whl/cu124')
return os.environ.get('TORCH_COMMAND', 'torch==2.4.1+cu124 torchvision==0.19.1+cu124 --index-url https://download.pytorch.org/whl/cu124')
def install_rocm_zluda():
@@ -561,7 +561,7 @@ def install_rocm_zluda():
if error is None:
try:
zluda_installer.load(zluda_path)
torch_command = os.environ.get('TORCH_COMMAND', f'torch=={zluda_installer.get_default_torch_version(device)}+cu118 torchvision --index-url https://download.pytorch.org/whl/cu118')
torch_command = os.environ.get('TORCH_COMMAND', f'torch=={zluda_installer.get_default_torch_version(device)} torchvision --index-url https://download.pytorch.org/whl/cu118')
log.info(f'Using ZLUDA in {zluda_path}')
except Exception as e:
error = e
@@ -636,7 +636,7 @@ def install_ipex(torch_command):
def install_openvino(torch_command):
check_python(supported_minors=[8, 9, 10, 11, 12], reason='OpenVINO backend requires Python 3.9, 3.10 or 3.11')
log.info('OpenVINO: selected')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.1+cpu torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cpu')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.1+cpu torchvision==0.18.1+cpu --index-url https://download.pytorch.org/whl/cpu')
install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2024.3.0'), 'openvino')
install(os.environ.get('ONNXRUNTIME_PACKAGE', 'onnxruntime-openvino'), 'onnxruntime-openvino', ignore=True)
install('nncf==2.12.0', 'nncf')
@@ -723,7 +723,7 @@ def check_torch():
torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
elif allow_directml and args.use_directml and ('arm' not in machine and 'aarch' not in machine):
log.info('DirectML: selected')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.4.1+cpu torchvision torch-directml')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.4.1 torchvision torch-directml')
if 'torch' in torch_command and not args.version:
install(torch_command, 'torch torchvision')
install('onnxruntime-directml', 'onnxruntime-directml', ignore=True)