Custom heads with Trainer API

Hi, I’ve been using this guide for fine-tuning a BERT-based model with the Trainer API:
My doubt is: Can I build a custom head on top of a model built with Trainer? For example, if I want to build my head with several layers for SQuAD, can I do it with Trainer? Or am I bound to use only the pre-existing heads? If it’s possible, do you have a snippet of code that shows how to do it? I couldn’t find any example.