torch训练简单例子

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

# 定义一个简单的模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

# 创建模型实例
model = MyModel()

# 检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将模型移动到GPU
model.to(device)

# 创建优化器
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 定义随机数生成的样本数量
batch_size = 32
# 定义训练轮数
num_epochs = 10

# 执行训练循环
for epoch in range(num_ep