Add pytorch nightly support for ROCm 6.0 (installer) (dev)

This commit is contained in:
Jonas Kanafani
2024-02-24 16:11:22 +01:00
committed by Vladimir Mandic
parent 5ac36c6bf1
commit 03ec491a69
+1 -1
View File
@@ -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}')