Changed code according to code review.