Encoder-Decoder Models Can't Generate using Apex
Author: ManavR123
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.4.2
- Platform: Linux-5.4.0-1041-aws-x86_64-with-debian-buster-sid
- Python version: 3.7.10
- PyTorch version (GPU?): 1.7.1+cu101 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
@patrickvonplaten , @patil-suraj 

## Information

Model I am using (Bert, XLNet ...): T5, ProphetNet

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

## To reproduce

``` Python
>>> from apex import amp
>>> from transformers import ProphetNetForConditionalGeneration, ProphetNetTokenizer
>>> tokenizer = ProphetNetTokenizer.from_pretrained("microsoft/prophetnet-large-uncased-squad-qg")
>>> model = ProphetNetForConditionalGeneration.from_pretrained("microsoft/prophetnet-large-uncased-squad-qg")
>>> model = model.to("cuda")
>>> model = amp.initialize(model, opt_level="O2") # comment out this line and it works fine
>>> encoder_inputs = tokenizer( ["Hello, I am"], return_tensors="pt", truncation=True, padding=True)["input_ids"].to("cuda")
>>> model.generate(encoder_inputs, num_beams=5, do_sample=True, max_length=32)
Traceback (most recent call last):
  File "ex.py", line 8, in <module>
    model.generate(encoder_inputs, num_beams=5, do_sample=True, max_length=32)
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/transformers/generation_utils.py", line 1093, in generate
    **model_kwargs,
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/transformers/generation_utils.py", line 1990, in beam_sample
    output_hidden_states=output_hidden_states,
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/apex/amp/_initialize.py", line 197, in new_fwd
    **applier(kwargs, input_caster))
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/transformers/models/prophetnet/modeling_prophetnet.py", line 1841, in forward
    return_dict=return_dict,
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/RLDiverseQG/env/lib/python3.7/site-packages/transformers/models/prophetnet/modeling_prophetnet.py", line 1732, in forward
    encoder_hidden_states=encoder_outputs[0],
KeyError: 0
```

You can switch out ProphetNet for T5ForConditionalGeneration and get the same error

## Expected behavior

I expect using apex shouldn't affect the code's functionality.

I figured out the main cause of the error: apex converts `BaseModelOutput` objects into dictionaries, but a lot of the code functionality relies on receiving the former. I don't know if there is a way to avoid this.

It is a pretty tedious fix to go over all of the places where this assumption is made and change direct indexing or attribute accesses to use `.get` but I believe that would be the solution to this problem. Hopefully, this is some helpful direction. I am also happy to help with this!

Comments:
I was able to get a fix working for ProphetNetForConditionalGeneration

Hi @ManavR123 

> I figured out the main cause of the error: apex converts BaseModelOutput objects into dictionaries, but a lot of the code functionality relies on receiving the former. I don't know if there is a way to avoid this.

You could pass `return_dict=False` to `forward` if you don't want the mode to return the output as model output classes, when `return_dict` is `False`, `tuple` is returned

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.