Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

clean up dead code

This commit is contained in:
William Falcon 2019-07-03 17:09:39 -04:00
parent 522af58504
commit 75e32daad4
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
pytorch_lightning/root_module/model_saving.py
View File

@ -2,7 +2,7 @@ import torch
import os
import re
import pdb
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
class ModelIO(object):
@ -66,7 +66,7 @@ class TrainerIO(object):
checkpoint['optimizer_states'] = optimizer_states
# request what to save from the model
model = self.model.module if type(self.model) is LightningDataParallel else self.model
model = self.model.module if type(self.model) is LightningDistributedDataParallel else self.model
checkpoint_dict = model.get_save_dict()
# merge trainer and model saving items