相关文章推荐
# ----------- 判断模型是在CPU还是GPU上 ---------------------- model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True) print(next(model.parameters()).device) # 输出:cpu model = model.cuda() print(next(model.parameters()).device) # 输出:cuda:0 model = model.cpu() print(next(model.parameters()).device) # 输出:cpu # ----------- 判断数据是在CPU还是GPU上 ---------------------- data = torch.ones([2, 3]) print(data.device) # 输出:cpu data = data.cuda() print(data.device) # 输出:cuda:0 data = data.cpu() print(data.device) # 输出:cpu

此外,用.is_cuda也可以判断模型和数据是否在GPU上,例如: data.is_cuda

模型和数据在GPU和CPU之间进行迁移的方法,更具体的可以参考: 把数据和模型迁移到GPU: .cuda()方法和.to(device)方法

 
推荐文章