重置用于交叉验证的网络[pytorch]的模型参数和权重

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

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

model = MyModel()

def reset_model(model):
    for layer in model.modules():
        if isinstance(layer, nn.Linear):
            layer.reset_parameters()

reset_model(model)

optimizer = optim.SGD(model.parameters(), lr=0.01)