默认是CPU，如果想要用GPU需要：

1. 安装配置cuda，然后更新/下载支持gpu版本的pytorch，可以参考：https://blog.csdn.net/weixin_35757704/article/details/124315569
2. 设置device：

   device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
   

   然后将数据与模型后面都额外加上.to(device)即可

示例程序

import torch
import torch.nn as nn# 一个简单的模型
class LinearRegressionModel(nn.Module):def __init__(self, input_shape, output_shape):super(LinearRegressionModel, self).__init__()self.linear = nn.Linear(input_shape, output_shape)def forward(self, x):out = self.linear(x)return outdef main():x_train = torch.randn(100, 4)  # 生成训练特征y_train = torch.randn(100, 1)  # 生成labelmodel = LinearRegressionModel(x_train.shape[1], 1)optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化函数criterion = nn.MSELoss()  # 损失函数for epoch in range(100):optimizer.zero_grad()outputs = model(x_train)loss = criterion(outputs, y_train)loss.backward()optimizer.step()if __name__ == '__main__':main()

修改为GPU版本：

import torch
import torch.nn as nn# 一个简单的模型
class LinearRegressionModel(nn.Module):def __init__(self, input_shape, output_shape):super(LinearRegressionModel, self).__init__()self.linear = nn.Linear(input_shape, output_shape)def forward(self, x):out = self.linear(x)return outdef main():# 1. 设置devicedevice = torch.device('cuda' if torch.cuda.is_available else 'cpu')# 2. 数据与模型后都加 .to(device) 即可x_train = torch.randn(100, 4).to(device)  # 生成训练特征y_train = torch.randn(100, 1).to(device)  # 生成labelmodel = LinearRegressionModel(x_train.shape[1], 1).to(device)  # next(transformer.parameters()).deviceoptimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化函数criterion = nn.MSELoss()  # 损失函数for epoch in range(100):optimizer.zero_grad()outputs = model(x_train)loss = criterion(outputs, y_train)loss.backward()optimizer.step()if __name__ == '__main__':main()

修改后：

1. 查看变量的位置：可以使用x_train.device查看tensor变量的位置
2. 查看模型的位置：可以使用next(model.parameters()).device查看模型的位置

注意：不在同一个位置上的变量之间无法计算，模型无法使用不在同一个位置的数据