Model.save_pretrained() does not save layer changes

Beginners
1

I trained a model based on BeitForImageClassification, and for some reason I needed to remove the classifier layer,I did it by calling this, and it worked fine.

pretrained_model.classifier = torch.nn.Identity()

And this is the current model structure:

========================================================================================================================
BeitForImageClassification                                             [1, 768]                  --
├─BeitModel: 1-1                                                       [1, 768]                  --
│    └─BeitEmbeddings: 2-1                                             [1, 197, 768]             768
│    │    └─BeitPatchEmbeddings: 3-1                                   [1, 196, 768]             590,592
│    │    └─Dropout: 3-2                                               [1, 197, 768]             --
│    └─BeitEncoder: 2-2                                                [1, 197, 768]             --
│    │    └─ModuleList: 3-3                                            --                        85,169,088
│    └─Identity: 2-3                                                   [1, 197, 768]             --
│    └─BeitPooler: 2-4                                                 [1, 768]                  --
│    │    └─LayerNorm: 3-4                                             [1, 768]                  1,536
├─Identity: 1-2                                                        [1, 768]                  --
========================================================================================================================

However it seems this change was not saved to the local model file with save_pretrain(), and when I called BeitForImageClassification.from_pretrained to read my saved model, the classifier went back to its previous state:

========================================================================================================================
BeitForImageClassification                                             [1, 589]                  --
├─BeitModel: 1-1                                                       [1, 768]                  --
│    └─BeitEmbeddings: 2-1                                             [1, 197, 768]             768
│    │    └─BeitPatchEmbeddings: 3-1                                   [1, 196, 768]             590,592
│    │    └─Dropout: 3-2                                               [1, 197, 768]             --
│    └─BeitEncoder: 2-2                                                [1, 197, 768]             --
│    │    └─ModuleList: 3-3                                            --                        85,169,088
│    └─Identity: 2-3                                                   [1, 197, 768]             --
│    └─BeitPooler: 2-4                                                 [1, 768]                  --
│    │    └─LayerNorm: 3-4                                             [1, 768]                  1,536
├─Linear: 1-2                                                          [1, 589]                  452,941
========================================================================================================================

Here’s my code:

new_model_path = model_path + "_new"
    pretrained_model = BeitForImageClassification.from_pretrained(model_path, ignore_mismatched_sizes=True)
    pretrained_model.classifier = torch.nn.Identity()
    summary(model=pretrained_model, input_size=(1, 3, 224, 224))
    pretrained_model.save_pretrained(new_model_path)

    new_model = BeitForImageClassification.from_pretrained(new_model_path, ignore_mismatched_sizes=True)
    summary(model=new_model, input_size=(1, 3, 224, 224))

Did I miss something?