I want to finetuning model by replacing only the last linear layer.
I could do that when I used DataParallel module as below
model = nn.DataParallel(model)

...
model.load_state_dict(checkpoint['state_dict'])
...    

model.module.fc = torch.nn.Linear(model.module.fc.in_features,
                                    opt.n_finetune_classes)
model.module.fc = model.module.fc.cuda()

In case of DDP, how can I do that?
according to the warning in the document 4,  it seems that I cannot change parameters after I load the checkpoint.
This module assumes all parameters are registered in the model by the time it is created. No parameters should be added nor removed later. Same applies to buffers.

Should I change the linear layer first and load the parameters without the linear layer?
Is that a right way?