I am trying to retrieve the probability distribution for each token predicted during beam search with multiple predictions (and multiple beams). For each prediction, I would like to get the probability distribution over tokens for each token predicted.

In the `generate`

function, I have included the `output_scores=True`

and `return_dict_in_generate=True`

keyword arguments. The relevant output is the scores, which is a list of tensors of shape `([num beams, num tokens])`

In the case with only 1 beam and 1 prediction, it is trivial to retrieve the probability distribution for each token: simply that beam. However, with multiple predictions and multiple beams, how can I retrieve the probability distribution over tokens for each token predicted for each prediction?

Much appreciated.