def topK(score):
    """
    Args: 
        score: [beam, num_vocab]
    Return:
        top_score: [beam]
        top_rowid: [beam], beam id 
        top_colid: [beam], word id
    """
    beam_size, num_vocab = score.size() 
    flat_score = score.view(beam_size * num_vocab)
    top_score, top_index = flat_score.topk(k=beam_size, dim=0)
    top_rowid, top_colid = top_index / num_vocab, top_index % num_vocab
    return top_score, top_rowid, top_colid