Map optimizer to correct device after loading from checkpoint.
Loading from `optimizer.pt` causes `optimizer` to be mapped to the same device as the saved `optimizer.pt`. In most cases it's `cuda:0`(saved by local master), which puts all optimizers on 
gpu0, causing OOM more easily in multi-gpu training.

Might fix issues like [#3730](https://github.com/huggingface/transformers/issues/3730). 
Thank you!