Upadted the code