checkpoint only on rank=0 now

2019-07-03 18:14:34 -04:00 · 2019-07-03 18:14:34 -04:00 · f101152650
parent 75e32daad4
commit f101152650
1 changed files with 2 additions and 1 deletions
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@ -588,4 +588,5 @@ class Trainer(TrainerIO):

        # model checkpointing
        print('save callback...')
-        self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
+        if self.proc_rank == 0:
+            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)