Change lightning module params to dict when loading (#1639)

* change module params to dict

* tiny change

* reverse
This commit is contained in:
Jacob Zhong
2020-05-02 08:44:35 -04:00
committed by GitHub
parent 210cd657dd
commit 4dc77b5a1a
+3 -2
View File
@@ -1550,8 +1550,9 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks):
)
# load the state_dict on the model automatically
model_args = [hparams] if hparams else []
model = cls(*model_args, *args, **kwargs)
if hparams:
kwargs.update(hparams=hparams)
model = cls(*args, **kwargs)
model.load_state_dict(checkpoint['state_dict'])
# give model a chance to load something