Fix cache for GPT-Neo-X
# What does this PR do?

As pointed out on #17745, there is a problem with the logic in the cache for GPT-Neo-X. Can confirm I can use it for generation after this PR, but not before.

Fixes #17745
Thanks a lot! If it's not too much work could you maybe try running this test: https://github.com/huggingface/transformers/blob/522a9ece4baeb5abfec8953ef76469a530e987d5/tests/models/gpt_neox/test_modeling_gpt_neox.py#L144 

I think right now it's not run (seems like the test function was removed)<|||||>_The documentation is not available anymore as the PR was closed or merged._<|||||>I can also do it also otherwise :-)<|||||>Added tests that were removed and had a corresponding `check` function. Thanks for flagging this!