From 0687e7b0203f1be7c70a11a1eccdc52b04b797f0 Mon Sep 17 00:00:00 2001 From: cmdr2 Date: Mon, 6 Jan 2025 19:48:44 +0530 Subject: [PATCH] Update the index url for AMD ROCm torch install --- scripts/check_modules.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/scripts/check_modules.py b/scripts/check_modules.py index 2cf40c46..d5c4eb05 100644 --- a/scripts/check_modules.py +++ b/scripts/check_modules.py @@ -92,16 +92,16 @@ def install(module_name: str, module_version: str): amd_gpus = setup_amd_environment() if module_name == "torch": if "Navi 3" in amd_gpus: - # No AMD 7x00 support in rocm 5.2, needs nightly 5.5. build - module_version = "2.1.0.dev-20230614+rocm5.5" - index_url = "https://download.pytorch.org/whl/nightly/rocm5.5" + # No AMD 7x00 support in rocm 5.2, needs 5.5+ + module_version = "2.1.0+rocm5.5" + index_url = "https://download.pytorch.org/whl/rocm6.2" else: module_version = "1.13.1+rocm5.2" elif module_name == "torchvision": if "Navi 3" in amd_gpus: - # No AMD 7x00 support in rocm 5.2, needs nightly 5.5. build - module_version = "0.16.0.dev-20230614+rocm5.5" - index_url = "https://download.pytorch.org/whl/nightly/rocm5.5" + # No AMD 7x00 support in rocm 5.2, needs 5.5+ + module_version = "0.16.0+rocm5.5" + index_url = "https://download.pytorch.org/whl/rocm6.2" else: module_version = "0.14.1+rocm5.2" elif os_name == "Darwin": @@ -285,7 +285,7 @@ def apply_torch_install_overrides(module_version: str): module_version += "+cu117" index_url = "https://download.pytorch.org/whl/cu117" elif is_amd_on_linux(): - index_url = "https://download.pytorch.org/whl/rocm5.2" + index_url = "https://download.pytorch.org/whl/rocm6.2" return module_version, index_url