Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Parse hip device id from arg
Browse files Browse the repository at this point in the history
Signed-off-by: aviator19941 <[email protected]>
aviator19941 committed Nov 25, 2024
1 parent 8479203 commit 9a8c72e
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion sharktank/sharktank/utils/export_artifacts.py
Original file line number Diff line number Diff line change
@@ -270,8 +270,9 @@ def iree_benchmark_vmfb(
f"--device=hip://{i}" for i in range(self.tensor_parallelism_size)
]
else:
hip_device_arg = int(hip_device_id.split("://")[1])
rocr_visible_devices = [
f"ROCR_VISIBLE_DEVICES={','.join(str(i) for i in range(int(hip_device_id) + 1))}"
f"ROCR_VISIBLE_DEVICES={','.join(str(i) for i in range(hip_device_arg + 1))}"
]
params = [f"--parameters=model={irpa_path}"]
devices = [f"--device={hip_device_id}"]

0 comments on commit 9a8c72e

Please sign in to comment.