[`Gpt-neo-x`] Fix gpt neo-x multi gpu training
# What does this PR do?

This PR attempts to solve some issues users may encounter when training gpt-neo-x in multiple GPUs

Related: https://github.com/lvwerra/trl/pull/210

This PR might be not needed so putting it as draft for now
_The documentation is not available anymore as the PR was closed or merged._<|||||>Yes that also what I have thought after opening the PR, I'll find a workaroud to set everything on the correct device