Minor bug fix and comments