Address reviewer's comments