Cleaned up the code for registering for notifications