How could I load a self-defined model?

Hi all,I define a simple model in terminal as following.

class simple(nn.Module):
def init(self):
super(simple, self).init()
self.net = nn.Linear(28 * 28, 2)
def forward(self, x):
x = np.array(x)
x = torch.from_numpy(x)
x = Variable(x)
x = x.view(1,1,28,28)
batch_size = x.size(0)
x = x.view(batch_size, -1)
output = self.net(x.float())
return out

model = simple()
And I save it by torch.save(model,‘simple.pkl’).

Then,I open a new terminal,but I cannot load the model above by torch.load(‘simple.pkl’). However,If I use a model from torchvision.models such as VGG,It works.Why this happens?How could I load my own model?

My recollection was that this should work - maybe it changed in new versions of pytorch? Which version are you using?

In any case, this is not the ideal way of tackling the problem. What if you modify the code of your class and want to load a model saved with the old code? It will be chaos.

So the best way is to only save the model state dict (which is the collection of the parameters of the model). So you would do something like

model = simple()
trainModel(model)
torch.save(model.state_dict(), 'my_trained_model_weights')

=== New session === 

model = simple()
model.load_state_dict(torch.load('my_trained_model_weights'))

Thanks.

I know this way,but is there any way in Pytorch that I can directly load a self-defined class model in new session without define the class in both session? It seems that if I don’t define the class in new session, and then load it, I will get AttributionError : No module named simple in this example.