Small fix and code clean up.