Save only part of model

#1

Hi I have a question about saving the model.
I currently doing some project regard to knowledge distillation, so I load the pretrained teacher model in the student model.
When I finish the training, only thing that I want to save is the parameters of student model, but it seems that teacher model also included in the model save file.

How can I remove those teacher model automatically? Seems that register_buffer do this job but cant find the details of how to use it for complex model

Thanks

#2

It depends on your currently used workflow and I assume you are storing the state_dict of a model containing both submodels. If so you might want to store the state_dict of the student model only by accessing it. If this description doesn’t fit your use case could you describe it a bit more?