Code Review comments fix