import torch


def _rearrange_output(logit, pred):
    """
    Rearrange the word logits so that the decoded word has the highest valued
    logits by swapping the indices predicted with those with maximum logits.
    """
    max_logits, max_logit_indices = torch.max(logit, 2, keepdim=True)
    pred_indices = pred.unsqueeze(2)
    pred_logits = torch.gather(logit, 2, pred_indices)
    logit_rearranged = logit.scatter(2, pred_indices, max_logits)
    logit_rearranged.scatter_(2, max_logit_indices, pred_logits)
    return logit_rearranged