AttributeError: Can‘t get attribute ‘Net‘ on module ‘__main__‘

When using the pytorch deep learning framework, when we load the pre-trained complete pkl model, if an error is reported:

AttributeError: Can't get attribute 'Net' on module '__main__'

At this point, we should load a training class in the main function to initialize the neural network, as follows:

# 加入：初始化
class Net(nn.Module):
    def __init__(self, hidden_layers=64):
        super(Net, self).__init__()
        # hidden_layers = 128
        self.fc1    =   nn.Linear(8, hidden_layers)
        self.fc2    =   nn.Linear(hidden_layers, hidden_layers)
        self.fc_mu  =   nn.Linear(hidden_layers, 1)
    def forward(self, x):
        x   =   F.relu(self.fc1(x))
        x   =   F.relu(self.fc2(x))
        mu  =   5*torch.sigmoid(self.fc_mu(x))
        return mu
# 加入：实例化
r_net = Net(hidden_layers=64)
# 加载预先训练模型的参数
r_net = torch.load('./model.pkl')      

At this point r_net can be used normally