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}')