Skip to content

Commit

Permalink
Added build.py change for ROCm EP (#811)
Browse files Browse the repository at this point in the history
Added option in `build.py` for ROCm
  • Loading branch information
skyline75489 authored Sep 9, 2024
1 parent de6a671 commit dac90d5
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/python/py/models/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
valid_gqa_configurations = [
("cpu", TensorProto.FLOAT),
("cuda", TensorProto.FLOAT16),
("rocm", TensorProto.FLOAT16),
("dml", TensorProto.FLOAT16),
]
if (self.ep, self.io_dtype) in valid_gqa_configurations:
Expand Down Expand Up @@ -2742,7 +2743,7 @@ def get_args():
"-e",
"--execution_provider",
required=True,
choices=["cpu", "cuda", "dml", "web"],
choices=["cpu", "cuda", "rocm", "dml", "web"],
help="Execution provider to target with precision of model (e.g. FP16 CUDA, INT4 CPU, INT4 WEB)",
)

Expand Down

0 comments on commit dac90d5

Please sign in to comment.