1 files changed, 4 insertions, 7 deletions
diff --git a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
index b710cf6bdb..55d789c66a 100644
--- a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
@@ -1,6 +1,8 @@
+# frozen_string_literal: true
+
 module ActiveRecord
-  # = Active Record Belongs To Polymorphic Association
   module Associations
+    # = Active Record Belongs To Polymorphic Association
     class BelongsToPolymorphicAssociation < BelongsToAssociation #:nodoc:
       def klass
         type = owner[reflection.foreign_type]
@@ -11,12 +13,7 @@ module ActiveRecord
 
         def replace_keys(record)
           super
-          owner[reflection.foreign_type] = record.class.base_class.name
-        end
-
-        def remove_keys
-          super
-          owner[reflection.foreign_type] = nil
+          owner[reflection.foreign_type] = record ? record.class.base_class.name : nil
         end
 
         def different_target?(record)