Address review comment