diff --git a/installer.py b/installer.py index 0602d9b37..9463fea81 100644 --- a/installer.py +++ b/installer.py @@ -699,7 +699,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.3.1 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)