Made changes as per code review feedback

#2396