Hi @lewtun !
I have almost finished my work on an extension of @VictorSanh  work, I tried to make it as generic as possible, to be able to patch any network with only minimal additional work, and to include it in your own training infrastructure.
As @VictorSanh mentionned, it won’t probably be part of transformers, but a standalone tool. I will be releasing it in the following weeks (hopefully before end of month), so I hope you can wait until that point !
To be able to patch a network “on the fly” you can use the approach I used in pytorch_block_sparse 7 , using inspection of pytorch modules.
(You can see the first results of  my work I had a few weeks ago here 6 for example )
François