Small changes to display code