Thanks a lot! Is there any general overview, how to work with multihead networks and custom regulatizations (e.g. add to loss something, which depends on model, not on model output)?
Upd.: seems like https://docs.fast.ai/callback.html#callback has everything I need.