Few changes in code