code changes done as per review comments
Changes after code review.