diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
index 86d2af1..bb7f521 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
@@ -99,7 +99,8 @@
      */
     void computeState(ITestDevice device, boolean visibleOnly, String[] waitForActivitiesVisible,
                       boolean compareTaskAndStackBounds) throws Exception {
-        waitForValidState(device, visibleOnly, waitForActivitiesVisible, null);
+        waitForValidState(device, visibleOnly, waitForActivitiesVisible, null,
+                compareTaskAndStackBounds);
 
         assertSanity();
         assertValidBounds(compareTaskAndStackBounds);
@@ -115,12 +116,15 @@
      *                 Pass null to skip this check.
      */
     void waitForValidState(ITestDevice device, boolean visibleOnly,
-                           String[] waitForActivitiesVisible, int[] stackIds) throws Exception {
+                           String[] waitForActivitiesVisible, int[] stackIds,
+                           boolean compareTaskAndStackBounds) throws Exception {
         int retriesLeft = 5;
         do {
+            // TODO: Get state of AM and WM at the same time to avoid mismatches caused by
+            // requesting dump in some intermediate state.
             mAmState.computeState(device);
             mWmState.computeState(device, visibleOnly);
-            if (shouldWaitForValidStacks()
+            if (shouldWaitForValidStacks(compareTaskAndStackBounds)
                     || shouldWaitForActivities(waitForActivitiesVisible, stackIds)) {
                 log("***Waiting for valid stacks and activities states...");
                 try {
@@ -153,7 +157,7 @@
         } while (retriesLeft-- > 0);
     }
 
-    private boolean shouldWaitForValidStacks() {
+    private boolean shouldWaitForValidStacks(boolean compareTaskAndStackBounds) {
         if (!taskListsInAmAndWmAreEqual()) {
             // We want to wait for equal task lists in AM and WM in case we caught them in the
             // middle of some state change operations.
@@ -166,6 +170,14 @@
             log("***stackBoundsInAMAndWMAreEqual=false");
             return true;
         }
+        try {
+            // Temporary fix to avoid catching intermediate state with different task bounds in AM
+            // and WM.
+            assertValidBounds(compareTaskAndStackBounds);
+        } catch (AssertionError e) {
+            log("***taskBoundsInAMAndWMAreEqual=false : " + e.getMessage());
+            return true;
+        }
         return false;
     }
 
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerActivityVisiblityTests.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerActivityVisiblityTests.java
index b485377..6f527ed 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerActivityVisiblityTests.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerActivityVisiblityTests.java
@@ -33,7 +33,7 @@
     public void testVisibleBehindHomeActivity() throws Exception {
         executeShellCommand(getAmStartCmd(VISIBLE_BEHIND_ACTIVITY));
         mAmWmState.waitForValidState(mDevice, true, new String[] {VISIBLE_BEHIND_ACTIVITY},
-                new int[] {FULLSCREEN_WORKSPACE_STACK_ID});
+                new int[] {FULLSCREEN_WORKSPACE_STACK_ID}, false /* compareTaskAndStackBounds */);
 
         executeShellCommand(AM_START_HOME_ACTIVITY_COMMAND);
         mAmWmState.waitForHomeActivityVisible(mDevice);
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java
index f148b4f..7a691ec 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerPinnedStackTests.java
@@ -52,7 +52,7 @@
         }
 
         mAmWmState.waitForValidState(mDevice, true, new String[] {topActivityName},
-                new int[] {PINNED_STACK_ID});
+                new int[] {PINNED_STACK_ID}, false /* compareTaskAndStackBounds */);
         mAmWmState.computeState(mDevice, null);
 
         if (supportsPip()) {