Map a list of list to TPU tensor
I have this code to deconstruct a list of list to different tensors.
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(torch.LongTensor, zip(*batch))

If I want to create these tensors on gpu, I can use below code:
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(torch.cuda.LongTensor, zip(*batch))

But now I want to create all these on TPU, what should I do? Is there anything like below?
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(torch.xla.LongTensor, zip(*batch))


You can use xla device following the guide here.
You can select the device and pass it to your function like this:
import torch_xla.core.xla_model as xm
device = xm.xla_device()
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(lambda x: torch.Tensor(x).to(device).long(), zip(*batch))

You can even parametrize the device variable, torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;) can be used to select between cuda and cpu.
