Minor corrections to code