Fix problems mentioned in the code review