UI bug fix