torch.Tensor.to的用法

更改数据类型torch.dtype或者设备torch.device

x = torch.Tensor([1, 2, 3])
print(x.dtype)  # torch.float32
x = x.to(torch.int64)
print(x.dtype)  # torch.int64

print(x.device)  # device(type='cpu')
print(x)  # tensor([1, 2, 3])
cuda0 = torch.device('cuda:0')
x = x.to(cuda0)
print(x.device)  # device(type='cuda', index=0)
print(x)  # tensor([1, 2, 3], device='cuda:0')

参考

torch.Tensor.to — PyTorch 1.12 documentation