Anyone knows of code for probing layers on QA tasks

Hello!

I would like to run an experiment and probe different layers of different LM for QA tasks. As in I want to train a model and then evaluate it using only x layers.
I was wondering if anyone knows of a package/code that would allow me to do that, so train a BERT model and then evaluate QA on different layers?

Thank you!