register_hook 1 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.