CUDA(Computer Unified Device Architecture),是NVIDIA推出的运算平台。CUDA是一种有NVIDIA推出的通用并行计算架构,该架构使GPU能够解决复杂的计算问题。
torch.cuda这个模块增加了对CUDA tensor的支持,能够在cpu和gpu上使用相同的方法操作tensor.
通过.to方法能够把一个tensor转移到另外一个设备(比如从CPU转到GPU)
import torch
x = torch.rand(5,3)
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.3356, 1.9082, 1.7105],
[1.2802, 1.1533, 1.3852],
[1.2288, 1.2393, 1.7443],
[1.3479, 1.6767, 1.8700],
[1.7983, 1.5379, 1.4596]], device='cuda:0')
tensor([[1.3356, 1.9082, 1.7105],
[1.2802, 1.1533, 1.3852],
[1.2288, 1.2393, 1.7443],
[1.3479, 1.6767, 1.8700],
[1.7983, 1.5379, 1.4596]], dtype=torch.float64)










