Map optimizer to correct device after loading from checkpoint.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 251f0dd4bc70e..25c0baeb14662 100644
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -331,7 +331,9 @@ def train(self, model_path: Optional[str] = None):
             and os.path.isfile(os.path.join(model_path, "scheduler.pt"))
         ):
             # Load in optimizer and scheduler states
-            optimizer.load_state_dict(torch.load(os.path.join(model_path, "optimizer.pt")))
+            optimizer.load_state_dict(
+                torch.load(os.path.join(model_path, "optimizer.pt"), map_location=self.args.device)
+            )
             scheduler.load_state_dict(torch.load(os.path.join(model_path, "scheduler.pt")))
 
         model = self.model

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). 