Multiple fixes in SageMakerTrainer
# What does this PR do?

This PR adds quite a few fixes to the `SageMakerTrainer` to make sure example scripts run fully. In particular it fixes:
- save made the training hanging forever
- predict didn't work
- evaluation required using `drop_last=True` which is not something anyone wants.

The goal is now to test a little bit more that functionality before merging the `SageMakerTrainer` into the main `Trainer` (otherwise one can't use model parallelism in seq2seq examples or QA example). The plan is to have them merged in the v4.5.0.