Slight changes to code