Address more review comments