Fix gradient checkpointing megatron bert
This PR fixes a bug that a user can encounter while using generate and models that use gradient_checkpointing.

Fixes Issue https://github.com/huggingface/transformers/issues/21737

cc @younesbelkada or @gante
_The documentation is not available anymore as the PR was closed or merged._