Small fixes after code review