Lightning isn`t automatically moving a tensor to certain device 
I am trying to implement this seq2seq cod from official pytorch docs  here 
I am using Lightning so i can scale the compute to TPU.
in a model there is an ouput tensor produced while decoding that collect a few samples of predictions on a single given input.
outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(device)      where i set device as device = xm.xla_device()
here i have to explicitly do   .to(device)  because Lightning cannot detect it.
But in TPU as computations are distributed across the TPU cores. hence even after moving the tensor to one core doesn`t help
i need a way to sync the tensor to all the tpu cores while the computation is at it.
how to sync the tensor to all TPU cores