Fixing issues from PR comments