Python-将模型和张量加载到GPU

在配置好GPU版本的pytorch的环境后，可以先检查一下cuda是否可用

import torch
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

或者是

torch.cuda.is_available()

 

然后有两种方式可以实现将模型和张量加载到GPU

model=model.to(device)
x=x.to(device)
y=y.to(device)

或者是

if torch.cuda.is_available():
  model = model.cuda()
  x = x.cuda()
  y = y.cuda()

推荐用第一种方法