small changes on code