diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index eaaf804..236982f 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -2652,6 +2652,12 @@
      */
     void setInputMethodWindowLocked(WindowState win) {
         mInputMethodWindow = win;
+        // Update display configuration for IME process.
+        if (mInputMethodWindow != null) {
+            final int imePid = mInputMethodWindow.mSession.mPid;
+            mService.mAtmInternal.onImeWindowSetOnDisplay(imePid,
+                    mInputMethodWindow.getDisplayId());
+        }
         computeImeTarget(true /* updateImeTarget */);
     }