Can I add hook when training and remove them when testing

Hello, I am training a model with evaluation periodically. So I wonder if I can add hook when training and remove then when evaluating. Could I use the hook handle all the time without registering new hook again?

register_hook as well as register_forward_hook will return a handle, which can be used to remove the hook via handle.remove(). So yes, you should be able to use the hook only during training and remove them afterwards.