pytorch cpu to gpu示例

import torch
import torch.nn as nn
import torch.optim as optim

#确保 GPU 是可用的
if torch.cuda.is_available():
    print("GPU is available")
else:
    print("GPU is not available")

#指定设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#定义模型
model = nn.Linear(5, 3)

#将模型移动到设备上
model.to(device)

#定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

#训练模型
for epoch in range(10):
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")