1 files changed, 5 insertions, 1 deletions
diff --git a/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java b/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java
index bb6ec6b1c..b2e979dba 100644
--- a/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java
+++ b/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java
@@ -402,7 +402,11 @@ public abstract class DynamicPredictionDictionaryBase extends ExpandableDictiona
     }
 
     public void registerUpdateSession(PersonalizationDictionaryUpdateSession session) {
-        session.setDictionary(this);
+        session.setPredictionDictionary(mLocale, this);
         mSessions.add(session);
     }
+
+    public void unRegisterUpdateSession(PersonalizationDictionaryUpdateSession session) {
+        mSessions.remove(session);
+    }
 }