Fixing removeClass function