Fixes related to comments on PR