Fixes after comments on pull request