This might be too late to help the OP, but for anyone who meets a similar problem:
“Taking an advantage of forward hook paradigm in PyTorch [30], torchdistill supports introducing such auxiliary modules without altering the original implementations
of the models.”
torchdistill: A Modular, Configuration-Driven Framework for Knowledge Distillation, Yoshitomo (2020 RRPR)
So it seems that torchdistill already has this functionality.
Even for those who are not directly dealing with knowledge distillation (both OP and me), it’s good to know that there is one solid implementation. This would be a nice starting point