Reload checkpoint
# What does this PR do?

This PR fixes a few bugs linked to resuming training from a checkpoint, mainly:
- the progress was not properly displayed (beginning at 0 instead of the step from the checkpoint)
- reloading the optimizer state and scheduler state on TPU was causing an error

Tested on TPU, single-GPU and multi-GPU env.

Fixes #4963 
Fixes #7976 
