Fix incorrect class registration