我有这段代码可以将一个列表分解为不同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之间进行选择。