diff --git a/mmdet/ops/nms/gpu_nms.pyx b/mmdet/ops/nms/gpu_nms.pyx index e5ae72578731c38150bf0c79866fcabfcb936ceb..65af87fd9a9377302ba82f24cffe114e8e10e730 100644 --- a/mmdet/ops/nms/gpu_nms.pyx +++ b/mmdet/ops/nms/gpu_nms.pyx @@ -19,7 +19,7 @@ memory_pool = {} def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, np.int32_t device_id=0): cdef int boxes_num = dets.shape[0] - cdef int boxes_dim = dets.shape[1] + cdef int boxes_dim = 5 cdef int num_out cdef size_t base cdef np.ndarray[np.int32_t, ndim=1] \ @@ -29,7 +29,7 @@ def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, cdef np.ndarray[np.int_t, ndim=1] \ order = scores.argsort()[::-1] cdef np.ndarray[np.float32_t, ndim=2] \ - sorted_dets = dets[order, :] + sorted_dets = dets[order, :5] cdef float cthresh = thresh if device_id not in memory_pool: with nogil: