Strange shape of Scores vector

Hello everyone,

I am using BART and I have enabled return_scores = True with beam_size of 4.
The shape of the Scores vector is (seq_len, batch_size * beam_size * num_returned sequence, vocab_size).

I would like to know how I can subdivide the vector obtaining the shape
(seq_len, batch_size, beam_size, num_returned_sequence, vocab_size).

Because at the moment it is not possible to map the input sentences in a batch with the output.

Kind regards,