Fine-tuning using our custom script?

[moha23]

Hi! I want to build my own model using just a few initial layers of some of the pertained models here. The following is how I'm doing it, and just want to make sure this is alright and not breaking anything internally? The available train scripts would require me to change my existing data loading procedures etc, this just saves time. (I save the weights separately and load manually in a different server where downloading from the internet is not allowed).

# saving weights for use later

model = timm.create_model('vit_base_patch16_224', pretrained=True)
torch.save(model.state_dict(), './vit-pretrained-weights.pt')

# on the target server

vit=timm.create_model('vit_large_patch16_224', img_size=256, pretrained=False)
vit.load_state_dict(torch.load('./vit-pretrained-weights.pt'))

class newModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = vit
        self.patch_embed = self.model.patch_embed
        self.pos_drop = self.model.pos_drop
        self.block = self.model.blocks[0]
        self.classifier = nn.Linear(1024, 1000)

    def forward(self, x):
        x = self.patch_embed(x)
        x = self.pos_drop(x)
        x = self.block(x)
        x = self.classifier(x)
        return x

model = newModel()

Train as usual using a criterion and optimizer reference code

What I expect is that the pertained weights are used for initialisation and then everything gets updated during training on new data. This seems to work fine (no error that is), but I'm not sure if I'm unknowingly breaking something. Are there any such requirements for using Timm models? Sorry if this is a naive question.

EDIT: I missed out on adding positional encoding so modified that.