pytorch Tensor在cpu和gpu之间转换

创建一个tensor

a = torch.arange(10)
'''
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
device(type='cpu') 默认在cpu上
'''

从cpu转到gpu上

a = torch.arange(10).cuda()
'''
device(type='cuda', index=0) 调用cuda()方法后Tensor存储在gpu
'''

从gpu转到cpu上

a = torch.arange(10).cuda()
'''
device(type='cuda', index=0) 调用cuda()方法后Tensor在gpu上计算
'''
a = a.cpu()
'''
device(type='cpu') 调用cpu()方法后Tensor又重新回到cpu
'''