Minor modifications to code