minor changes to code