Add sleep to fix the rendezous error

Author: kaushikb11

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -277,7 +277,11 @@ def _wrapped_function(
         result = function(*args, **kwargs)
         if self.local_rank == 0:
             return_queue.put(move_data_to_device(result, "cpu"))
+
         self.barrier("end-process")
+        # https://github.com/pytorch/xla/issues/2190#issuecomment-641665358
+        if self.local_rank == 0:
+            time.sleep(2)
 
     def _worker_setup(self, process_idx: int):
         reset_seed()