updates

2019-09-27 23:35:04 +02:00 · 2019-09-27 23:35:04 +02:00 · b421afa508
parent c6d3efbf95
commit b421afa508
1 changed files with 1 additions and 1 deletions
--- a/utils/utils.py
+++ b/utils/utils.py
@ -479,7 +479,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):

        # Select predicted classes
        class_conf = class_conf[i]
-        class_pred = class_pred[i].unsqueeze(1).float()
+        class_pred = class_pred[i].unsqueeze(1).type_as(class_conf)

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        pred[:, :4] = xywh2xyxy(pred[:, :4])