Add pytorch-triton-rocm to dependencies when targeting AMD because accelerate needs to find it in the rocm repo

2026-02-01 01:00:20 +08:00 · 2024-05-22 21:21:23 -07:00 · 2024-05-22 21:21:23 -07:00 · c0fc1d1458
commit c0fc1d1458
parent 0fcd07962f
1 changed files with 1 additions and 0 deletions
--- a/setup.py
+++ b/setup.py
@ -146,6 +146,7 @@ def dependencies(force_nightly: bool = False) -> List[str]:
        index_urls += [nvidia_torch_index]
    elif _is_amd():
        index_urls += [amd_torch_index]
+        _dependencies += ["pytorch-triton-rocm"]
    else:
        index_urls += [cpu_torch_index]