As i got from the concept of captioning that i need the features for the images not to classify it. So if i have a classification model, the classification head can be dropped. does that code doing the concept ?  the code here
the modifications are

def gradient(model, freeze: bool):
transformer = VisionTransformer(num_classes=0)
for parameter in transformer.parameters():
parameter.requires_grad_(not freeze)


def vit_small(patch_size=16, **kwargs):
model = VisionTransformer(         patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,         qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)         gradient(model, freeze=True)
return model

Does that right ?
Does after that i need re -train the model for my datasets after the modifications ?
What should result expected ?