FP16 NMS (#814)
Browse files- utils/general.py +0 -2
utils/general.py
CHANGED
@@ -596,8 +596,6 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
|
|
596 |
Returns:
|
597 |
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
598 |
"""
|
599 |
-
if prediction.dtype is torch.float16:
|
600 |
-
prediction = prediction.float() # to FP32
|
601 |
|
602 |
nc = prediction[0].shape[1] - 5 # number of classes
|
603 |
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
596 |
Returns:
|
597 |
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
598 |
"""
|
|
|
|
|
599 |
|
600 |
nc = prediction[0].shape[1] - 5 # number of classes
|
601 |
xc = prediction[..., 4] > conf_thres # candidates
|