Function to convert model output into a word

word_from_output.py

# Define a function which converts output into word
def wordFromOutput(output):
    '''
    Functions returns an index from the vocabulary and the corresponding word
    '''
    top_n, top_i = output.topk(1)
    category_i = top_i[0].item()
    return [k for (k, v) in vocab.items() if v == category_i], category_i