diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 9aca17b8fc..1105a6009f 100644
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -389,7 +389,9 @@ class Trainer:
             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