diff --git a/core/java/android/window/ImeOnBackInvokedDispatcher.java b/core/java/android/window/ImeOnBackInvokedDispatcher.java
index 9152e7837b828..a0bd7f70ca585 100644
--- a/core/java/android/window/ImeOnBackInvokedDispatcher.java
+++ b/core/java/android/window/ImeOnBackInvokedDispatcher.java
@@ -220,16 +220,14 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
      * @param previous the previously focused {@link ViewRootImpl}.
      * @param current the currently focused {@link ViewRootImpl}.
      */
-    // TODO(b/232845902): Add CTS to test IME back behavior when there's root view change while
-    // IME is up.
     public void switchRootView(ViewRootImpl previous, ViewRootImpl current) {
         for (ImeOnBackInvokedCallback imeCallback : mImeCallbacks) {
             if (previous != null) {
                 previous.getOnBackInvokedDispatcher().unregisterOnBackInvokedCallback(imeCallback);
             }
             if (current != null) {
-                current.getOnBackInvokedDispatcher().registerOnBackInvokedCallback(
-                        imeCallback.mPriority, imeCallback);
+                current.getOnBackInvokedDispatcher().registerOnBackInvokedCallbackUnchecked(
+                        imeCallback, imeCallback.mPriority);
             }
         }
     }