Revert "Fix to prevent class being removed"

This reverts commit f9374018978279bd24572d65a3b89dee0c5a5d1a.