Minor improvements

Let me know if you need me to make any changes to the PR.