Few code fixes