wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

fixed hpc save, load. cleaned apu

Browse Source
This commit is contained in:
William Falcon
2019-07-26 22:07:02 -04:00
parent 265411572f
commit 64de447545
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/root_module/root_module.py
+3 -4
View File
@@ -108,13 +108,12 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
else:
checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
# load the state_dict on the model automatically
model = cls(hparams)
model.load_state_dict(checkpoint['state_dict'])
# give model a chance to load something
model.load_model_specific(checkpoint)
# load the state_dict on the model automatically
model.load_state_dict(checkpoint['state_dict'])
model.on_load_checkpoint(checkpoint)
return model