slight changes to code