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

@@ -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 = (

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()

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: