diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 5a6851e..37e4902 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -4364,7 +4364,6 @@
         // Reset the last saved PiP snap fraction on removal.
         mDisplayContent.mPinnedTaskController.onActivityHidden(mActivityComponent);
         mDisplayContent.onRunningActivityChanged();
-        mWmService.mEmbeddedWindowController.onActivityRemoved(this);
         mRemovingFromDisplay = false;
     }
 
diff --git a/services/core/java/com/android/server/wm/EmbeddedWindowController.java b/services/core/java/com/android/server/wm/EmbeddedWindowController.java
index 98027bb..c9bae12 100644
--- a/services/core/java/com/android/server/wm/EmbeddedWindowController.java
+++ b/services/core/java/com/android/server/wm/EmbeddedWindowController.java
@@ -135,19 +135,6 @@
         return mWindowsByWindowToken.get(windowToken);
     }
 
-    void onActivityRemoved(ActivityRecord activityRecord) {
-        for (int i = mWindows.size() - 1; i >= 0; i--) {
-            final EmbeddedWindow window = mWindows.valueAt(i);
-            if (window.mHostActivityRecord == activityRecord) {
-                final WindowProcessController processController =
-                        mAtmService.getProcessController(window.mOwnerPid, window.mOwnerUid);
-                if (processController != null) {
-                    processController.removeHostActivity(activityRecord);
-                }
-            }
-        }
-    }
-
     static class EmbeddedWindow implements InputTarget {
         final IWindow mClient;
         @Nullable final WindowState mHostWindowState;
@@ -230,6 +217,13 @@
                 mInputChannel.dispose();
                 mInputChannel = null;
             }
+            if (mHostActivityRecord != null) {
+                final WindowProcessController wpc =
+                        mWmService.mAtmService.getProcessController(mOwnerPid, mOwnerUid);
+                if (wpc != null) {
+                    wpc.removeHostActivity(mHostActivityRecord);
+                }
+            }
         }
 
         @Override