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

finished dist (#911)

Browse Source
This commit is contained in:
William Falcon
2020-02-21 20:39:12 -05:00
committed by GitHub
parent 6e7dc9c236
commit c00a8a10dd
3 changed files with 33 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/trainer.py
+5
View File
@@ -983,6 +983,11 @@ class Trainer(TrainerIOMixin,
if self.use_ddp or self.use_ddp2:
dist.barrier()
# wait for all models to restore weights
if self.on_tpu and XLA_AVAILABLE:
# wait for all processes to catch up
torch_xla.core.xla_model.rendezvous()
# set up checkpoint callback
self.configure_checkpoint_callback()