Updating UI code to fix past commit