Revert "Revert "Fix to prevent class being removed""

This reverts commit 297e9f7e088e2fff10fa26790121df3486bd46cb.