Which hidden states have the highest score in beam search?

Hello, I’m currently utilizing the transformers library for inference purposes. I have enabled output_hidden_states by setting it to True , and I’ve specified a num_beams value of 4 . Consequently, the shape of hidden_states is [n_beams, n_tokens] . However, I’m uncertain about identifying the specific beam from n_beams that correlates with the highest-scoring inference result.

Here is an example:

hidden_states.shape == [4, 50]
output_tokens.shape == [50]
# hidden_states[0] <--> output_tokens ?
# hidden_states[1] <--> output_tokens ?
# hidden_states[2] <--> output_tokens ?
# hidden_states[3] <--> output_tokens ?