Skip to content

Commit 5612778

Browse files
loadamsgyou2021
authored andcommitted
Update setup.py handling of ROCm cupy (deepspeedai#7051)
Signed-off-by: gyou2021 <ganmei.you@intel.com>
1 parent 3ca3e2f commit 5612778

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

setup.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,8 @@ def get_env_if_set(key, default: typing.Any = ""):
102102
cupy = None
103103
if is_rocm_pytorch:
104104
rocm_major, rocm_minor = rocm_version
105-
# XXX cupy support for rocm 5 is not available yet.
106-
if rocm_major <= 4:
105+
# cupy support for rocm>5.0 is not available yet.
106+
if (rocm_major == 5 and rocm_minor == 0) or rocm_major <= 4:
107107
cupy = f"cupy-rocm-{rocm_major}-{rocm_minor}"
108108
else:
109109
cuda_major_ver, cuda_minor_ver = installed_cuda_version()

0 commit comments

Comments
 (0)