Hub device mismatch bug fix (#1594)
This commit is contained in:
parent
4a07280884
commit
e285034b4b
@ -265,7 +265,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None,
|
||||
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
||||
"""
|
||||
|
||||
nc = prediction[0].shape[1] - 5 # number of classes
|
||||
nc = prediction.shape[2] - 5 # number of classes
|
||||
xc = prediction[..., 4] > conf_thres # candidates
|
||||
|
||||
# Settings
|
||||
@ -277,7 +277,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None,
|
||||
merge = False # use merge-NMS
|
||||
|
||||
t = time.time()
|
||||
output = [torch.zeros(0, 6)] * prediction.shape[0]
|
||||
output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
|
||||
for xi, x in enumerate(prediction): # image index, image inference
|
||||
# Apply constraints
|
||||
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user