Some fixes after code review