Small tweaks to code