Fix few bugs