Fixed code review comments