Minor changes to code