1 files changed, 13 insertions, 0 deletions
diff --git a/lib/sqlalchemy/orm/identity.py b/lib/sqlalchemy/orm/identity.py
index 0fa541194..24dd47859 100644
--- a/lib/sqlalchemy/orm/identity.py
+++ b/lib/sqlalchemy/orm/identity.py
@@ -187,6 +187,12 @@ class WeakInstanceDict(IdentityMap):
             return list(self._dict.values())
 
     def discard(self, state):
+        st = self._dict.pop(state.key, None)
+        if st:
+            assert st is state
+            self._manage_removed_state(state)
+
+    def safe_discard(self, state):
         if state.key in self._dict:
             st = self._dict[state.key]
             if st is state:
@@ -259,6 +265,13 @@ class StrongInstanceDict(IdentityMap):
         state._instance_dict = self._wr
 
     def discard(self, state):
+        obj = self._dict.pop(state.key, None)
+        if obj is not None:
+            self._manage_removed_state(state)
+            st = attributes.instance_state(obj)
+            assert st is state
+
+    def safe_discard(self, state):
         if state.key in self._dict:
             obj = self._dict[state.key]
             st = attributes.instance_state(obj)