Add rocm to transformers/benchmark.py script (#68)

Co-authored-by: Stefan Sokolovic <[email protected]>
ROCm · Oct 2, 2024 · 8495918 · 8495918
1 parent 9c64ac0
commit 8495918
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py
@@ -788,7 +788,7 @@ def main():
         logger.error("fp16 is for GPU only")
         return
 
-    if args.precision == Precision.INT8 and args.use_gpu and args.provider != "migraphx":
+    if args.precision == Precision.INT8 and args.use_gpu and args.provider not in ["migraphx", "rocm"]:
         logger.error("int8 is for CPU only")
         return