Fixing few of the bugs