// Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred)
detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)
// Iterate through all predicted classes
for c in detections[:, -1].unique():
// Get the detections with the particular class
detections_class = detections[detections[:, -1] == c]
// Sort the detections by maximum objectness confidence
After Change
// Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred)
detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)
// Iterate through all predicted classes
unique_labels = detections[:, -1].cpu().unique()if prediction.is_cuda:
unique_labels = unique_labels.cuda()
for c in unique_labels:
// Get the detections with the particular class
detections_class = detections[detections[:, -1] == c]
// Sort the detections by maximum objectness confidence
