Finetuning Pix2struct with custom Image, QnA form

Hi all,

I have started fine-tuning the pix2struct model using custom dataset, leveraging the loss from the models directly. Is there a way we can freeze and unfreeze (requires grad function) for a portion of the model and enable grad-fn just for the head?
Also there is an attribute called the head masks here, what is the utility of that?