Added changes per code review