Cannot use torch.save(model.state_dict(), PATH) with transfer learning

1

I keep getting this error:

for param_tensor in model.state_dict():
    AttributeError: 'NoneType' object has no attribute 'state_dict'

Here is a code snippet from the transfer learning I am trying to implement:

...
model = torchvision.models.densenet161(pretrained=True)

for param in model.parameters():
    param.requires_grad = False

num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.classifier.parameters(), lr=learning_rate, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

model = train_model(model, criterion, optimizer, exp_lr_scheduler, num_epochs=num_epochs)

torch.save(model.state_dict(), './save.pt')

The model is able to get trained and output results, but it is the torch.save(model.state_dict(), PATH) part that does not work - any ideas??

2 
model = train_model(model, criterion, optimizer, exp_lr_scheduler, num_epochs=num_epochs)

What does function train_model return ? Are you sure it returns the model ? Because your code seems right.

3

I am copying the example from here:
https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html

the function was supposed to return the model, but I had somehow removed that line. Not that it is back, torch.save is working fine! Thank you for making me look there!

1 Like