diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 748557b..63a6960 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -175,8 +175,7 @@
                 && mPrefs.getBoolean(SHARED_PREFS_STASHED_KEY, DEFAULT_STASHED_PREF);
         boolean isInSetup = !mActivity.isUserSetupComplete() || sharedState.setupUIVisible;
         updateStateForFlag(FLAG_STASHED_IN_APP_MANUAL, isManuallyStashedInApp);
-        // TODO(b/204384193): Temporarily disable SUW specific logic
-        // updateStateForFlag(FLAG_STASHED_IN_APP_SETUP, isInSetup);
+        updateStateForFlag(FLAG_STASHED_IN_APP_SETUP, isInSetup);
         if (isInSetup) {
             // Update the in-app state to ensure isStashed() reflects right state during SUW
             updateStateForFlag(FLAG_IN_APP, true);