diff --git a/installer.py b/installer.py index 0c25f0446..2d6ffab15 100644 --- a/installer.py +++ b/installer.py @@ -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 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 --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)} 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)}+cu118 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 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 --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 torchvision torch-directml') + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.4.1+cpu torchvision torch-directml') if 'torch' in torch_command and not args.version: install(torch_command, 'torch torchvision') install('onnxruntime-directml', 'onnxruntime-directml', ignore=True)