我有这段代码可以将一个列表分解为不同Tensor。
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(torch.LongTensor, zip(*batch))
如果我想在gpu上创建这些Tensor,我可以使用下面的代码:
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(torch.cuda.LongTensor, zip(*batch))
但是现在我想在TPU上创建所有这些,我应该怎么做?有没有像下面这样的东西?
token_a_index, token_b_index, isNext, input_ids, segment_ids, masked_tokens, masked_pos = map(torch.xla.LongTensor, zip(*batch))
1条答案
按热度按时间wydwbb8l1#
您可以按照here指南使用xla设备。
您可以选择设备并将其传递给函数,如下所示:
您甚至可以参数化设备变量,
torch.device("cuda" if torch.cuda.is_available() else "cpu")
可用于在cuda
和cpu
之间进行选择。