Small fix and upgrade