This task is about classification.
I want the model to make predictions with priority information, not just probability, when predicting a true label.
For example, there are 6 classes and I have the priority information like this  [0,0,0,1,1,2], threshold: 0.5
Lower number means more important class!
If the difference in probability between the label with more important priority and the current prediction label not exceed the threshold, I change the current prediction label to a label with more important priority.
For example, probability from models [0.3, 0.2, 0.15, 0.15, 0.1, 0.1] and priority [1,0,1,1,1,1], threshold:0.2.
Normally, the predicted class should be 0th, but it is changed to 1th according to priority information.
(1th class is more important class and 0.3-0.2<0.2)
I implemented this logic with the for loop below, but I think it can be optimized with matrix operation. How can I do it?
def modify_predResult_with_priority(prob, priority, threshold):
    #prob: [batch_size, num_cls]
    #priority: [num_cls] ex) 6class: [0,0,0,1,1,2]
    #threshold: 0.5
   
    sorted_prob, sorted_prob_cls_index = prob.sort(descending=True)
    sorted_prob_priority = torch.tensor(priority)[sorted_prob_cls_index]

    candidate_cls = sorted_prob_cls_index[:,0]
    candidate_prob = sorted_prob[:,0]
    candidate_priority = sorted_prob_priority[:,0]

    for batch in range(sorted_prob.shape[0]):
        for idx in range(sorted_prob.shape[1]):
            if idx==0:
                continue
            if ((candidate_prob[batch] - sorted_prob[batch][idx]) > threshold):
                break
            if (sorted_prob_priority[batch][idx] < candidate_priority[batch]):
                candidate_prob[batch] = sorted_prob[batch][idx]
                candidate_cls[batch] = sorted_prob_cls_index[batch][idx]
                candidate_priority[batch] = sorted_prob_priority[batch][idx]
    return candidate_cls