minor updates to the code