import torch import numpy as np device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") x = torch.tensor(np.arange(15).reshape(3,5)) if torch.cuda.is_available(): device = torch.device("cuda") y = torch.ones_like(x,device=device) x = x.to(device) z = x + y print(z) print(z.to("cpu",torch.double))
tensor([[ 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10],
[11, 12, 13, 14, 15]], device='cuda:0', dtype=torch.int32)
tensor([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.],
[11., 12., 13., 14., 15.]], dtype=torch.float64)
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:pytorch中CUDA类型的转换 - Python技术站