Change the hook to make it work without core modifications