import torch


def _rearrange_output(logit, pred):
    """
    Rearrange the word logits so that the decoded word has the highest valued
    logits.
    """
    for batch_idx, v_path in enumerate(pred):
        for w_idx, word in enumerate(v_path):
            w_logits = logit[batch_idx][w_idx]
            v_label = word.item()
            # make the word on the optimal path the greatest
            _, maxIndex = torch.max(w_logits, 0)
            w_logits[v_label], w_logits[maxIndex] = (
                w_logits[maxIndex].item(),
                w_logits[v_label].item(),
            )
            logit[batch_idx][w_idx] = w_logits
    return logit