SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Map optimizer to correct device after loading from checkpoint. (#4403)

Browse Source 
* Map optimizer to correct device after loading from checkpoint.

* Make style test pass

Co-authored-by: Julien Chaumond <chaumond@gmail.com>
This commit is contained in:
Shaoyen
2020-05-18 20:16:05 -07:00
committed by GitHub
parent bf14ef75f1
commit 384f0eb2f9
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/trainer.py
View File

@@ -389,7 +389,9 @@ class Trainer:
and os.path.isfile(os.path.join(model_path, "scheduler.pt")) and os.path.isfile(os.path.join(model_path, "scheduler.pt"))
): ):
# Load in optimizer and scheduler states # 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"))) scheduler.load_state_dict(torch.load(os.path.join(model_path, "scheduler.pt")))
model = self.model model = self.model