Fixes based on review comments