for current_epoch in range(all_epoch): model.train() for idx, (train_x, train_label) in enumerate(train_loader): train_x = train_x.to(device) train_label = train_label.to(device) sgd.zero_grad() predict_y = model(train_x.float()) loss = loss_fn(predict_y, train_label.long()) loss.backward() sgd.step() print(f'current epoch:{current_epoch}')

这段代码是一个简单的训练模型的代码示例。每个 epoch 都会对训练集进行一次完整的遍历。在每个 epoch 中，模型会被设置为训练模式（model.train()），然后遍历训练集的每个样本。对于每个样本，我们首先将输入数据和标签数据移动到设备上（通常是 GPU），然后将优化器的梯度清零（sgd.zero_grad()）。接下来，我们使用模型进行预测（model(train_x.float())），计算预测结果与真实标签之间的损失（loss_fn(predict_y, train_label.long())），并根据损失计算梯度（loss.backward()）。最后，我们使用优化器更新模型的参数（sgd.step()）。在每个 epoch 结束时，打印当前 epoch 的信息。

这段代码可以用来训练分类任务的模型，你可以根据自己的数据和模型进行相应的修改。