diff --git a/lib/set_as_primary.rb b/lib/set_as_primary.rb
index 3d5b863..0250762 100644
--- a/lib/set_as_primary.rb
+++ b/lib/set_as_primary.rb
@@ -92,6 +92,6 @@ def _polymorphic_condition_options
     end
 
     def _klass
-      self.class
+      self.class.base_class
     end
 end