From d7264ed2c3be6d0ae9a9c1c1055e91c9e22f06dc Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Mon, 30 Sep 2024 12:46:55 +0900 Subject: [PATCH] upgrade torch for directml backend --- installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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)