Fixing highlight bug