Cleaned up custom module code a bit