Fixing UI issues