Minor bug fix

Minor bug fixes