Made small changes to code