diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 87dffc828fc60..7931cd58d977b 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -197,6 +197,7 @@ import android.view.ViewRootImpl;
 import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.view.WindowManagerPolicyConstants.PointerEventListener;
+import android.window.ITaskOrganizer;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.MetricsLogger;
@@ -3425,10 +3426,7 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
 
     private void setInputMethodTarget(WindowState target, boolean targetWaitingAnim) {
         // Always update control target. This is needed to handle rotation.
-        // We cannot set target as the control target, because mInputMethodTarget can only help
-        // decide the z-order of IME, but cannot control IME. Only the IME target reported from
-        // updateInputMethodTargetWindow can control IME.
-        updateImeControlTarget(mInputMethodControlTarget);
+        updateImeControlTarget(target);
         if (target == mInputMethodTarget && mInputMethodTargetWaitingAnim == targetWaitingAnim) {
             return;
         }