Support rocm device for R50 and Retinanet (mlcommons#1730)

Co-authored-by: Miro <[email protected]> Co-authored-by: Pablo Gonzalez <[email protected]>
GATEOverflow · Jun 27, 2024 · 9e2c9f6 · 9e2c9f6
1 parent 763883b
commit 9e2c9f6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vision/classification_and_detection/run_common.sh b/vision/classification_and_detection/run_common.sh
@@ -19,7 +19,7 @@ device="cpu"
 for i in $* ; do
     case $i in
        tf|onnxruntime|tflite|pytorch|tvm-onnx|tvm-pytorch|tvm-tflite|ncnn) backend=$i; shift;;
-       cpu|gpu) device=$i; shift;;
+       cpu|gpu|rocm) device=$i; shift;;
        gpu) device=gpu; shift;;
        resnet50|mobilenet|ssd-mobilenet|ssd-resnet34|ssd-resnet34-tf|retinanet) model=$i; shift;;
     esac