Fix gradient checkpointing bug in gptneox
# What does this PR do?

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._<|||||>Apologies, hand't pushed the latest commit. All done now!<|||||>Awesome, thank you for the contribution <3 