Remove layer in pretrained model

I’m new to ML, I’m trying to perform an ablation study. Is there a concise way to access and remove layers from a pretrained model? Specifically, I’m working on this code:

pretrained_model_name = "nvidia/mit-b0"
model = AutoModelForSemanticSegmentation.from_pretrained(
    pretrained_model_name, id2label=id2label, label2id=label2id
)

# to-do: something like "model.remove_layers(-1)"

I know there’s model.config, but it doesn’t seem to have what I’m looking for. Any tips? Thanks!

1 Like