Fix issue removing custom class