Addressed comments on the pull request