import torch


def convert_output_to_tensor(output, vocab):
    """
    Converts the latent dimension into data again the opposite of data_process
    This version also accounts for the probability outputs using torch.argmax()
    """
    imm = [
        [vocab.get_itos()[torch.argmax(prob).item()] for prob in result]
        for result in output
    ]
    ret = []
    for row in imm:
        new_row = []
        for i in row:
            if i == "<pad>":
                new_row.append(0)
            elif i == "<unk>":
                new_row.append(-1)
            else:
                new_row.append(int(i))
        ret.append(new_row)

    return torch.tensor(ret)