Possibility to monitor head training?

[Ulipenitz]

Hello!
Is there a plan to implement monitoring of the head training phase as well?
Currently it seems that logging (Wandb, neptune etc.) stops after the body training, but it can be useful to see the logs & implement early stopping for the head as well.
Especially for end-to-end training when you use more complex pytorch heads where you modify the architecture during hp-search.