点此进入CSDN

点此添加QQ好友 加载失败时会显示




你的浏览器不支持播放哦!!nuttertools 您的浏览器不支持该播放!

pytorch中CUDA类型的转换

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)

  

posted @ 2020-02-12 16:11  高颜值的殺生丸  阅读(1488)  评论(0)    收藏  举报

作者信息

昵称:

刘新宇

园龄:4年6个月


粉丝:1209


QQ:522414928