ROCm 6.0 lock mpmath to 1.2.1

This commit is contained in:
Disty0
2024-03-21 00:26:14 +03:00
parent cbd4941e95
commit e2cacf8732
+4 -1
View File
@@ -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}')