Updated code for production