Update for review comments