From e2cacf87324c814f2e32bce3cbd41e6bba226329 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 21 Mar 2024 00:26:14 +0300 Subject: [PATCH] ROCm 6.0 lock mpmath to 1.2.1 --- installer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/installer.py b/installer.py index a01e5d4ca..de70a12da 100644 --- a/installer.py +++ b/installer.py @@ -525,7 +525,10 @@ 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", "6.0"}: + if rocm_ver == "6.0": + # mpmath 1.3.0 crashes at import + torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision mpmath==1.2.1 --pre --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}') + elif rocm_ver == "5.7": 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}')