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

set dp as default backend

Browse Source
This commit is contained in:
William Falcon
2019-07-18 11:45:55 -04:00
parent 3a1525222d
commit 6d1d5ef68e
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/root_module/model_saving.py
+3 -2
View File
@@ -2,7 +2,7 @@ import torch
import os
import re
import pdb
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
class ModelIO(object):
@@ -66,7 +66,8 @@ class TrainerIO(object):
checkpoint['optimizer_states'] = optimizer_states
# request what to save from the model
model = self.model.module if type(self.model) is LightningDistributedDataParallel else self.model
is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
model = self.model.module if is_dp_module else self.model
checkpoint_dict = model.get_save_dict()
# merge trainer and model saving items