Great, hope it works out!
A quick note is that the callback expects the model to be a sequential model, it looks for sub-modules [x in model] where x.parameters() has length greater than one by default. Alternatively you can pass (with partial) a list of the modules you want to register hooks on explicitly.