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

Fixes Issue #21737 

cc @sgugger @amyeroberts
_The documentation is not available anymore as the PR was closed or merged._<|||||>@sgugger @amyeroberts As mentioned I did the following:-

> Thanks for your PR! You need to remove the other one below (line 660).

However won't that cause an issue with function declaration and it's corresponding else blocks? I saw other implementations of this fix and they don't remove the block below just add it again above. What do you think?<|||||>Hi @krypticmouse - thanks for your question and applying the update @sgugger requested. 

Looking at the diff again, L654 shouldn't be removed. I believe this should resolve the function declaration issue you mentioned. <|||||>You will also need to resolved the conflict as `logger.warning` has been renamed to `logger.warning_once` since you opened your PR.<|||||>Is this ok to be merged now?