Minor updates to notebook