Bug fix in UI