Minor updates on code