diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
index 5489455..dda50ca 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
@@ -167,8 +167,10 @@
 
     /**
      * Indicates that the IME window has re-parented to the new target when the IME control changed.
+     *
+     * @param displayId the display hosting the IME window
      */
-    public abstract void onImeParentChanged();
+    public abstract void onImeParentChanged(int displayId);
 
     /**
      * Destroys the IME surface for the given display.
@@ -300,7 +302,7 @@
                 }
 
                 @Override
-                public void onImeParentChanged() {
+                public void onImeParentChanged(int displayId) {
                 }
 
                 @Override
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index c3eb519..0d29b7d 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -5671,7 +5671,7 @@
         }
 
         @Override
-        public void onImeParentChanged() {
+        public void onImeParentChanged(int displayId) {
             synchronized (ImfLock.class) {
                 // Hide the IME method menu only when the IME surface parent is changed by the
                 // input target changed, in case seeing the dialog dismiss flickering during
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index ba22763..c98280e 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -4794,7 +4794,8 @@
             assignRelativeLayerForIme(getSyncTransaction(), true /* forceUpdate */);
             scheduleAnimation();
 
-            mWmService.mH.post(() -> InputMethodManagerInternal.get().onImeParentChanged());
+            mWmService.mH.post(
+                    () -> InputMethodManagerInternal.get().onImeParentChanged(getDisplayId()));
         } else if (mImeControlTarget != null && mImeControlTarget == mImeLayeringTarget) {
             // Even if the IME surface parent is not changed, the layer target belonging to the
             // parent may have changes. Then attempt to reassign if the IME control target is