From 5778a4131c84fd7ed424dfd294acf3f1f2779273 Mon Sep 17 00:00:00 2001 From: srush Date: Sun, 23 Feb 2020 15:00:32 -0500 Subject: [PATCH] Add tags to the rendezvous calls for TPU. (#921) * Update data_loading.py * Update training_io.py * Update trainer.py --- pytorch_lightning/trainer/data_loading.py | 2 +- pytorch_lightning/trainer/trainer.py | 2 +- pytorch_lightning/trainer/training_io.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pytorch_lightning/trainer/data_loading.py b/pytorch_lightning/trainer/data_loading.py index ee545128..43e87928 100644 --- a/pytorch_lightning/trainer/data_loading.py +++ b/pytorch_lightning/trainer/data_loading.py @@ -233,7 +233,7 @@ class TrainerDataLoadingMixin(ABC): self.get_val_dataloaders() # wait for all processes to catch up - torch_xla.core.xla_model.rendezvous() + torch_xla.core.xla_model.rendezvous("pl.TrainerDataLoadingMixin.get_dataloaders") # support IterableDataset for train data self.is_iterable_train_dataloader = ( diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 37c59a11..4d204bd2 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -1002,7 +1002,7 @@ class Trainer(TrainerIOMixin, # 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() + torch_xla.core.xla_model.rendezvous("pl.Trainer.run_pretrain_routine") # set up checkpoint callback self.configure_checkpoint_callback() diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py index 1bc14966..569c5718 100644 --- a/pytorch_lightning/trainer/training_io.py +++ b/pytorch_lightning/trainer/training_io.py @@ -183,7 +183,7 @@ class TrainerIOMixin(ABC): # 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() + torch_xla.core.xla_model.rendezvous("pl.TrainerIOMixin.restore_weights") # clear cache after restore if self.on_gpu: