small changes to code