举一个使用这个函数的例子

好的，下面是一个使用上述训练函数的例子：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载数据
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,))])
train_set = datasets.MNIST('../data', train=True, download=True,
                           transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

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

# 训练模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
num_epochs = 10
trained_model = train(model, train_loader, criterion, optimizer, device, num_epochs)

# 保存模型
torch.save(trained_model.state_dict(), 'trained_model.pth')

在上述例子中，我们定义了一个简单的神经网络模型，用于 MNIST 数据集的分类任务。然后我们加载数据集，定义损失函数和优化器，并使用上述训练函数 `train` 进行训练。最后，我们将训练好的模型保存到文件中。