当前位置:网站首页>pytorch总结—TENSOR ON GPU

pytorch总结—TENSOR ON GPU

2022-06-29 08:31:00 TJMtaotao

用方法 to() 可以将 Tensor 在CPU和GPU(需要硬件⽀支持)之间相互移动。

if torch.cuda.is_available(): #返回True or False #当GPU可用为True,反之则为False
device = torch.device("cuda") # GPU
y = torch.ones_like(x, device=device) # 直接创建⼀个在GPU上的Tensor
x = x.to(device) # 等价于 .to("cuda")

z = x + y
print(z)
print(z.to("cpu", torch.double)) # to()还可以同时更改数据类型

 

原网站

版权声明
本文为[TJMtaotao]所创,转载请带上原文链接,感谢
https://blog.csdn.net/TJMtaotao/article/details/103763497