

TITLE
Measuring time when using xla_spawn on multiple cores 

COMMENTS
2

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
Hi,
I would like to measure the time of the program, when this is launched in distributed way on mutiple TPU cores with xla_spawn, I tried to write some callback like this which retunrs the same results with num_cores = 1 and 8, I am not sure how to handle it when this is distributed? I want to make sure pytorch XLA is using all cores indeed, thank you for your help. 

```
class PrinterTimeCallback(TrainerCallback):
  
    #A bare :class:`~transformers.TrainerCallback` that just prints the logs.
   
    def on_epoch_begin(self, args, state, control, **kwargs):
        """
        Event called at the beginning of an epoch.
        """
        if state.is_local_process_zero:
            self.start_time = time.time()


    def on_epoch_end(self, args, state, control, **kwargs):
        """
        Event called at the end of an epoch.
        """
        if state.is_local_process_zero:
            self.end_time = time.time()
            total_time = self.end_time - self.start_time
            print("total_time ", total_time)

```
   
