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

load_spawn_weights only in proc rank 0 (#1385) 

Co-authored-by: Alexander Reshytko <areshytko@Alexanders-MacBook-Pro.local>
This commit is contained in:
areshytko 2020-04-06 17:17:16 +03:00 committed by GitHub
parent 4ed3027309
commit 9754c5da55
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
pytorch_lightning/trainer/distrib_data_parallel.py
View File

@ -363,6 +363,10 @@ class TrainerDDPMixin(ABC):
:param model: :param model:
:return: :return:
""" """
loaded_model = original_model
if self.proc_rank == 0:
# load weights saved in ddp # load weights saved in ddp
path = os.path.join(self.default_save_path, '__temp_weight_ddp_end.ckpt') path = os.path.join(self.default_save_path, '__temp_weight_ddp_end.ckpt')
loaded_model = original_model.__class__.load_from_checkpoint(path) loaded_model = original_model.__class__.load_from_checkpoint(path)