num_beams error in GPT2DoubleHead model

[vibhavagarwal5]

Environment info

• transformers version: 2.9.1
• Platform: Linux
• Python version: 3.6
• PyTorch version (GPU?): 1.5
• Tensorflow version (GPU?):
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: Yes

Who can help

@LysandreJik @patil-suraj

Information

I am trying to use model.generate() for the GPT2DoubleHeadModel but the beam search is giving an error.
Setting the num_beams > 1 results in the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/hdd1/vibhav/anaconda3/envs/vesnli/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/home/hdd1/vibhav/anaconda3/envs/vesnli/lib/python3.7/site-packages/transformers/modeling_utils.py", line 1125, in generate
    model_specific_kwargs=model_specific_kwargs,
  File "/home/hdd1/vibhav/anaconda3/envs/vesnli/lib/python3.7/site-packages/transformers/modeling_utils.py", line 1481, in _generate_beam_search
    past = self._reorder_cache(past, beam_idx)
  File "/home/hdd1/vibhav/anaconda3/envs/vesnli/lib/python3.7/site-packages/transformers/modeling_utils.py", line 1551, in _reorder_cache
    return tuple(layer_past.index_select(1, beam_idx) for layer_past in past)
  File "/home/hdd1/vibhav/anaconda3/envs/vesnli/lib/python3.7/site-packages/transformers/modeling_utils.py", line 1551, in <genexpr>
    return tuple(layer_past.index_select(1, beam_idx) for layer_past in past)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

However, things are working fine for num_beams=1 and for GPT2LMHeadModel(both beam search and non beam search)

[adamlin120]

encountered the same issue

[patil-suraj]

I think @patrickvonplaten might have some ideas.