Use association_class method which returns the reflection class, this method is redefined in polymorphic belongs to associations.

activerecord/lib/active_record/associations/association_proxy.rb

@@ -254,8 +254,8 @@ module ActiveRecord
           return nil unless defined?(@loaded)
 
           if !loaded? and (@owner.persisted? || foreign_key_present)
-            if IdentityMap.enabled? && ActiveRecord.const_defined?(@reflection.class_name)
-              @target = IdentityMap.get(@reflection.klass, @owner[@reflection.association_foreign_key])
+            if IdentityMap.enabled? && association_class
+              @target = IdentityMap.get(association_class, @owner[@reflection.association_foreign_key])
             end
             @target ||= find_target
           end
@@ -313,6 +313,10 @@ module ActiveRecord
         def we_can_set_the_inverse_on_this?(record)
           false
         end
+
+        def association_class
+          @reflection.klass
+        end
     end
   end
 end