From 03ec491a694fbb8f34f98d7dead41869c2dcc939 Mon Sep 17 00:00:00 2001 From: Jonas Kanafani Date: Sat, 24 Feb 2024 16:11:22 +0100 Subject: [PATCH] Add pytorch nightly support for ROCm 6.0 (installer) (dev) --- installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/installer.py b/installer.py index 1ee471bd3..2b2f80321 100644 --- a/installer.py +++ b/installer.py @@ -488,7 +488,7 @@ def check_torch(): log.info('Using CPU-only torch') torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision') else: - if rocm_ver in {"5.7"}: + if rocm_ver in {"5.7", "6.0"}: torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}') elif rocm_ver in {"5.5", "5.6"}: torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}')