Skip to content

Commit 11c35cd

Browse files
author
Rahul Batra
committed
[ROCm]: Add rocm as platform in export_harnessess_multi_platform_test
1 parent 82ee60b commit 11c35cd

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

tests/export_harnesses_multi_platform_test.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,9 +136,12 @@ def export_and_compare_to_native(
136136
if d.platform not in unimplemented_platforms
137137
]
138138
logging.info("Using devices %s", [str(d) for d in devices])
139-
# lowering_platforms uses "cuda" instead of "gpu"
139+
# lowering_platforms uses "cuda" or "rocm" instead of "gpu"
140+
gpu_platform = "cuda"
141+
if jtu.is_device_rocm():
142+
gpu_platform = "rocm"
140143
lowering_platforms: list[str] = [
141-
p if p != "gpu" else "cuda"
144+
p if p != "gpu" else gpu_platform
142145
for p in ("cpu", "gpu", "tpu")
143146
if p not in unimplemented_platforms
144147
]

0 commit comments

Comments
 (0)