Fix to sorting code