Fix loading the best model on the last stage of training
# What does this PR do?

It fixes the best model loading on the last stage of the training.

Fixes #11666

## Who can review?
@sgugger 


Great! Lat thing is to run `make style` on your branch to make the CI pass. Let me know if you run into any issue, I can also push on your branch.<|||||>> Great! Lat thing is to run `make style` on your branch to make the CI pass. Let me know if you run into any issue, I can also push on your branch.

CI is already fixed<|||||>Thanks again for the fix!