1. 首页
  2. 数据库
  3. 其它
  4. Pytorch to(device)用法

Pytorch to(device)用法

上传者: 2021-01-01 02:27:58上传 PDF文件 54.4KB 热度 20次
如下所示: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device) 这两行代码放在读取数据之前。 mytensor = my_tensor.to(device) 这行代码的意思是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去,之后的运算都在GPU上进行。 这句话需要写的次数等于需要保存GPU上的tensor变量的个数;一般情况下这些tensor变量都是最开始读数据时的tensor变量,后面衍生的变量自然也都在GPU上 如果是多个GP
下载地址
用户评论