Minor updates to code