Fix gradient checkpointing tests for `encoder-decoder` models
# What does this PR do?

The recently added tests (#18697) didn't send the model to the correct device and caused some test failed. This PR fixes it.
_The documentation is not available anymore as the PR was closed or merged._<|||||>Sorry about that! Thanks for fixing @ydshieh 