Error while generating more then one Beam output in T5

Intermediate
1

Hi everyone, I am using T5 model to generate more than one beam output using the beam search algorithm. When I set the number_of_return_sequences more than one the trainer return the following error:

Traceback (most recent call last):
  File "./pipeline.py", line 531, in <module>
    main()
  File "./pipeline.py", line 499, in main
    num_return_sequences=data_args.num_return_sequences, 
  File "/raid/kaushal/mxnlg/xl-sum/seq2seq_PAS/transformers/trainer_seq2seq.py", line 153, in predict
    return super().predict(test_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/raid/kaushal/mxnlg/xl-sum/seq2seq_PAS/transformers/trainer.py", line 1407, in predict
    test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
  File "/raid/kaushal/mxnlg/xl-sum/seq2seq_PAS/transformers/trainer.py", line 1497, in prediction_loop
    preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
  File "/raid/kaushal/mxnlg/xl-sum/seq2seq_PAS/transformers/trainer_pt_utils.py", line 331, in add_arrays
    slice_len = self._nested_set_tensors(self._storage, arrays)
  File "/raid/kaushal/mxnlg/xl-sum/seq2seq_PAS/transformers/trainer_pt_utils.py", line 350, in _nested_set_tensors
    i * slice_len : (i + 1) * slice_len
ValueError: could not broadcast input array from shape (800,512) into shape (100,512)

Any suggestions will really useful. @valhalla @sshleifer @patrickvonplaten Thank you!