Finetuning pretrained pytorch model

finetune.py

# from : http://pytorch.org/docs/master/notes/autograd.html
# pretrained models here : https://github.com/pytorch/vision/tree/master/torchvision/models
model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
# Replace the last fully-connected layer
# Parameters of newly constructed modules have requires_grad=True by default
model.fc = nn.Linear(512, 100)

# Optimize only the classifier
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)