akamaster / pytorch_resnet_cifar10

Proper implementation of ResNet-s for CIFAR10/100 in pytorch that matches description of the original paper.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

How to save the model?

haibao-yu opened this issue · comments

Hi,

When I set the BasicBlock's option as "A" and save the model with following code:
torch.save(model, 'test.pkl'),
here is the error:
"AttributeError: Can't pickle local object 'BasicBlock.init..'"

What's problem?

Thanks!

it is not recommended to save model, you should save model's state instead:
torch.save(model.state_dict(), 'test.pkl")