From c4c0ee8fc35937cfa940fdaaaf6b9660f5b355f5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 26 Sep 2022 14:13:03 +0200 Subject: [PATCH] Apple MPS -> CPU NMS fallback strategy (#9600) Until more ops are fully supported this update will allow for seamless MPS inference (but slower MPS to CPU transfer before NMS, so slower NMS times). Partially resolves https://github.com/ultralytics/yolov5/issues/9596 Signed-off-by: Glenn Jocher Signed-off-by: Glenn Jocher --- utils/general.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/utils/general.py b/utils/general.py index de7871c..a855691 100644 --- a/utils/general.py +++ b/utils/general.py @@ -843,6 +843,8 @@ def non_max_suppression( if isinstance(prediction, (list, tuple)): # YOLOv5 model in validation model, output = (inference_out, loss_out) prediction = prediction[0] # select only inference output + if 'mps' in prediction.device.type: # MPS not fully supported yet, convert tensors to CPU before NMS + prediction = prediction.cpu() bs = prediction.shape[0] # batch size nc = prediction.shape[2] - nm - 5 # number of classes xc = prediction[..., 4] > conf_thres # candidates