Small changes to code